company-ai-platform/tests/test_personal_data_erasure.py

import json
from collections.abc import Iterator
from datetime import timedelta
from uuid import uuid4

import pytest
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from sqlalchemy import create_engine, event, select
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.application.feishu.personal_data import (
    PERSONAL_DATA_ERASURE_ACTION,
    FeishuPersonalDataService,
)
from app.core.config import get_settings
from app.core.database import Base, get_db
from app.core.utils.time import utc_now
from app.modules.ai_memory.models import AIMemoryEntry
from app.modules.audit.models import AuditLog
from app.modules.business.models import MarketWatchlist
from app.modules.feishu_users.constants import FeishuUserRole
from app.modules.feishu_users.models import (
    FeishuAdminBootstrapTombstone,
    FeishuUser,
)
from app.modules.feishu_users.principal import FeishuPrincipal
from app.modules.feishu_users.routes import router
from app.modules.feishu_users.services import FeishuIdentityService
from app.modules.personalization.models import (
    AIConversation,
    AIConversationMessage,
    PersonalDataErasureRequest,
    UserPreference,
)
from app.modules.personalization.schemas import ErasureResult
from app.modules.subscriptions.models import PushDelivery, PushSubscription


@pytest.fixture
def session_factory() -> Iterator[sessionmaker[Session]]:
    engine = create_engine(
        "sqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )

    @event.listens_for(engine, "connect")
    def enable_foreign_keys(dbapi_connection, _connection_record) -> None:
        cursor = dbapi_connection.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()

    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    try:
        yield factory
    finally:
        Base.metadata.drop_all(engine)
        engine.dispose()


def _create_user(
    db: Session,
    suffix: str,
    *,
    role: str = FeishuUserRole.USER,
) -> FeishuUser:
    user = FeishuUser(
        code=f"FSU-{suffix}",
        tenant_key=f"tenant-{suffix}",
        open_id=f"ou-open-{suffix}",
        union_id=f"on-union-{suffix}",
        user_id=f"cli-user-{suffix}",
        role=role,
        timezone="Asia/Shanghai",
    )
    db.add(user)
    db.flush()
    return user


def _seed_personal_data(db: Session, user: FeishuUser) -> None:
    db.add(
        UserPreference(
            owner_id=user.id,
            category="tone",
            value="简洁",
            normalized_value="简洁",
            source="explicit",
        )
    )
    conversation = AIConversation(
        owner_id=user.id,
        chat_type="private",
        chat_key=user.open_id,
    )
    db.add(conversation)
    db.flush()
    db.add(
        AIConversationMessage(
            conversation_id=conversation.id,
            role="user",
            content="仅属于我的历史",
        )
    )
    db.add(
        AIMemoryEntry(
            code=f"MEM-{uuid4().hex}",
            owner_id=user.id,
            kind="memory",
            scope="user",
            subject="preference",
            content="个人记忆",
            source="explicit",
            actor=user.open_id,
        )
    )
    db.add(
        MarketWatchlist(
            owner_id=user.id,
            actor=user.open_id,
            symbol="600000.SH",
            enabled=True,
        )
    )
    subscription = PushSubscription(
        code=f"SUB-{uuid4().hex}",
        owner_id=user.id,
        target_type="user",
        target_id=user.open_id,
        prompt="个人提醒",
        schedule_type="daily",
        schedule_config={"hour": 9, "minute": 0},
        timezone="Asia/Shanghai",
        next_run_at=utc_now() + timedelta(days=1),
        status="active",
        consented_at=utc_now(),
    )
    db.add(subscription)
    db.flush()
    db.add(
        PushDelivery(
            code=f"DEL-{uuid4().hex}",
            subscription_id=subscription.id,
            scheduled_for=utc_now(),
            idempotency_key=uuid4().hex,
            message_uuid=str(uuid4()),
            status="pending",
            next_attempt_at=utc_now(),
        )
    )
    db.add(
        AuditLog(
            actor=f"feishu:{user.open_id}",
            source="feishu",
            action="test.personal.action",
            target_type="feishu-user",
            target_id=user.code,
            request_payload=json.dumps(
                {
                    "open_id": user.open_id,
                    "union_id": user.union_id,
                    "private_note": "仅属于该用户的请求内容",
                }
            ),
            response_payload=json.dumps(
                {
                    "user_id": user.user_id,
                    "code": user.code,
                    "private_answer": "仅属于该用户的响应内容",
                }
            ),
            request_id=f"request-{user.code}",
        )
    )
    db.commit()


def test_confirmation_is_one_user_only_and_expires(
    session_factory: sessionmaker[Session],
) -> None:
    with session_factory() as db:
        first = _create_user(db, "first")
        second = _create_user(db, "second")
        db.commit()
        first_principal = FeishuPrincipal.from_user(first)
        second_principal = FeishuPrincipal.from_user(second)
        service = FeishuPersonalDataService(db)

        confirmation = service.request_confirmation(first_principal)
        with pytest.raises(HTTPException) as wrong_code:
            service.confirm(first_principal, "WRONG")
        assert wrong_code.value.status_code == 422

        with pytest.raises(HTTPException) as cross_user:
            service.confirm(second_principal, confirmation.confirmation_code)
        assert cross_user.value.status_code == 422

        request = db.execute(
            select(PersonalDataErasureRequest).where(
                PersonalDataErasureRequest.owner_id == first.id
            )
        ).scalar_one()
        request.expires_at = utc_now() - timedelta(seconds=1)
        db.commit()
        with pytest.raises(HTTPException) as expired:
            service.confirm(first_principal, confirmation.confirmation_code)
        assert expired.value.status_code == 422
        assert db.get(FeishuUser, first.id) is not None
        assert db.get(FeishuUser, second.id) is not None


def test_confirm_erases_all_personal_data_and_anonymizes_audit(
    session_factory: sessionmaker[Session],
) -> None:
    with session_factory() as db:
        user = _create_user(db, "erase")
        other = _create_user(db, "other")
        _seed_personal_data(db, user)
        _seed_personal_data(db, other)
        owner_id = user.id
        tenant_key = user.tenant_key
        open_id = user.open_id
        identifiers = (user.code, user.open_id, user.union_id, user.user_id)
        principal = FeishuPrincipal.from_user(user)
        service = FeishuPersonalDataService(db)

        confirmation = service.request_confirmation(principal)
        result = service.confirm(principal, confirmation.confirmation_code)

        assert result.deleted["deliveries"] == 1
        assert result.deleted["subscriptions"] == 1
        assert result.deleted["preferences"] == 1
        assert result.deleted["conversation_messages"] == 1
        assert result.deleted["conversations"] == 1
        assert result.deleted["ai_memory"] == 1
        assert result.deleted["watchlist"] == 1
        assert result.deleted["identity"] == 1
        assert db.get(FeishuUser, owner_id) is None
        assert db.get(FeishuUser, other.id) is not None
        assert db.scalar(
            select(PushSubscription).where(PushSubscription.owner_id == owner_id)
        ) is None
        assert db.scalar(
            select(UserPreference).where(UserPreference.owner_id == owner_id)
        ) is None
        assert db.scalar(
            select(AIMemoryEntry).where(AIMemoryEntry.owner_id == owner_id)
        ) is None
        assert db.scalar(
            select(MarketWatchlist).where(MarketWatchlist.owner_id == owner_id)
        ) is None
        assert db.scalar(
            select(AIConversation).where(AIConversation.owner_id == owner_id)
        ) is None

        logs = list(db.execute(select(AuditLog)).scalars())
        serialized_logs = "\n".join(
            " ".join(
                str(value or "")
                for value in (
                    log.actor,
                    log.target_id,
                    log.request_payload,
                    log.response_payload,
                )
            )
            for log in logs
        )
        assert all(identifier not in serialized_logs for identifier in identifiers)
        final_log = db.execute(
            select(AuditLog).where(AuditLog.action == PERSONAL_DATA_ERASURE_ACTION)
        ).scalar_one()
        assert final_log.actor == result.anonymous_id
        assert final_log.target_type is None
        assert final_log.target_id is None
        assert final_log.request_payload is None
        assert final_log.response_payload is None
        assert final_log.request_id is None
        assert final_log.status == "success"

        anonymized_history = db.execute(
            select(AuditLog).where(
                AuditLog.action == "test.personal.action",
                AuditLog.actor == result.anonymous_id,
            )
        ).scalar_one()
        assert anonymized_history.target_type is None
        assert anonymized_history.target_id is None
        assert anonymized_history.request_payload is None
        assert anonymized_history.response_payload is None
        assert anonymized_history.request_id is None
        assert anonymized_history.status == "success"

        recreated = FeishuIdentityService(db).resolve_or_register(
            tenant_key=tenant_key,
            open_id=open_id,
        )
        assert recreated.owner_id != owner_id
        assert recreated.user_code not in identifiers
        assert recreated.role == FeishuUserRole.USER


def test_last_active_admin_is_protected_until_another_admin_exists(
    session_factory: sessionmaker[Session],
) -> None:
    with session_factory() as db:
        admin = _create_user(db, "admin", role=FeishuUserRole.ADMIN)
        db.commit()
        principal = FeishuPrincipal.from_user(admin)
        service = FeishuPersonalDataService(db)
        confirmation = service.request_confirmation(principal)

        with pytest.raises(HTTPException) as protected:
            service.confirm(principal, confirmation.confirmation_code)
        assert protected.value.status_code == 409
        assert db.get(FeishuUser, admin.id) is not None

        second_admin = _create_user(
            db,
            "second-admin",
            role=FeishuUserRole.ADMIN,
        )
        db.commit()
        result = service.confirm(principal, confirmation.confirmation_code)
        assert result.deleted["identity"] == 1
        assert db.get(FeishuUser, admin.id) is None
        assert db.get(FeishuUser, second_admin.id) is not None


def test_erased_configured_admin_is_not_bootstrapped_again(
    session_factory: sessionmaker[Session],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv(
        "FEISHU_ADMIN_IDENTITIES",
        "tenant-bootstrap:ou-bootstrap",
    )
    get_settings.cache_clear()
    try:
        with session_factory() as db:
            identity = FeishuIdentityService(db)
            initial = identity.resolve_or_register(
                tenant_key="tenant-bootstrap",
                open_id="ou-bootstrap",
            )
            assert initial.role == FeishuUserRole.ADMIN
            _create_user(db, "replacement-admin", role=FeishuUserRole.ADMIN)
            db.commit()

            service = FeishuPersonalDataService(db)
            confirmation = service.request_confirmation(initial)
            result = service.confirm(initial, confirmation.confirmation_code)

            tombstone = db.scalar(select(FeishuAdminBootstrapTombstone))
            assert tombstone is not None
            assert len(tombstone.identity_hash) == 64
            assert "tenant-bootstrap" not in tombstone.identity_hash
            assert "ou-bootstrap" not in tombstone.identity_hash

            recreated = FeishuIdentityService(db).resolve_or_register(
                tenant_key="tenant-bootstrap",
                open_id="ou-bootstrap",
            )
            assert recreated.owner_id != initial.owner_id
            assert recreated.role == FeishuUserRole.USER
            assert result.deleted["identity"] == 1
    finally:
        get_settings.cache_clear()


def test_extra_hook_failure_rolls_back_the_deletion_transaction(
    session_factory: sessionmaker[Session],
) -> None:
    with session_factory() as db:
        user = _create_user(db, "rollback")
        _seed_personal_data(db, user)
        owner_id = user.id
        principal = FeishuPrincipal.from_user(user)
        confirmation = FeishuPersonalDataService(db).request_confirmation(principal)

        def fail_hook(_db: Session, _owner_id: int, _anonymous_id: str) -> None:
            raise RuntimeError("integration erasure failed")

        with pytest.raises(RuntimeError, match="integration erasure failed"):
            FeishuPersonalDataService(
                db,
                extra_hooks=(fail_hook,),
            ).confirm(principal, confirmation.confirmation_code)

        assert db.get(FeishuUser, owner_id) is not None
        assert db.scalar(
            select(UserPreference).where(UserPreference.owner_id == owner_id)
        ) is not None
        assert db.scalar(
            select(AIMemoryEntry).where(AIMemoryEntry.owner_id == owner_id)
        ) is not None
        assert db.scalar(
            select(PushSubscription).where(PushSubscription.owner_id == owner_id)
        ) is not None
        assert db.scalar(
            select(PersonalDataErasureRequest).where(
                PersonalDataErasureRequest.owner_id == owner_id
            )
        ) is not None


def test_delete_api_uses_service_principal_actor_and_ignores_body_identity(
    session_factory: sessionmaker[Session],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("API_KEY", "service-key")
    monkeypatch.setenv("API_ACTOR", "service-admin")
    monkeypatch.setenv("API_KEYS", "[]")
    get_settings.cache_clear()
    seen: dict[str, str] = {}

    def fake_erase(
        _service: FeishuPersonalDataService,
        code: str,
        *,
        actor: str,
    ) -> ErasureResult:
        seen.update(code=code, actor=actor)
        return ErasureResult(
            anonymous_id="anonymous-test",
            deleted={"identity": 1},
        )

    monkeypatch.setattr(
        FeishuPersonalDataService,
        "erase_by_user_code",
        fake_erase,
    )
    app = FastAPI()
    app.include_router(router, prefix="/api/v1/integrations/feishu")

    def override_db() -> Iterator[Session]:
        with session_factory() as db:
            yield db

    app.dependency_overrides[get_db] = override_db
    client = TestClient(app)
    try:
        path = "/api/v1/integrations/feishu/users/FSU-target/personal-data"
        assert client.delete(path).status_code == 401
        assert (
            client.delete(path, headers={"X-API-Key": "wrong-key"}).status_code
            == 401
        )
        response = client.request(
            "DELETE",
            path,
            headers={"X-API-Key": "service-key"},
            json={
                "actor": "forged-user",
                "open_id": "ou-forged",
            },
        )
        assert response.status_code == 200
        assert response.json()["anonymous_id"] == "anonymous-test"
        assert seen == {
            "code": "FSU-target",
            "actor": "service-admin",
        }
    finally:
        get_settings.cache_clear()