company-ai-platform/tests/test_personalization_cleanup.py

from collections.abc import Iterator
from datetime import timedelta

import pytest
from fastapi import FastAPI
from sqlalchemy import create_engine, func, select
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from app.application.scheduling import create_scheduler
from app.core.config import get_settings
from app.core.database import Base
from app.core.utils.time import utc_now
from app.modules.feishu_users.models import FeishuUser
from app.modules.personalization.models import AIConversation, AIConversationMessage
from app.modules.personalization.services import ConversationService


@pytest.fixture(autouse=True)
def _reset_settings() -> Iterator[None]:
    get_settings.cache_clear()
    yield
    get_settings.cache_clear()


@pytest.fixture
def session_factory() -> Iterator[sessionmaker[Session]]:
    engine = create_engine(
        "sqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    Base.metadata.create_all(engine)
    factory = sessionmaker(bind=engine, expire_on_commit=False)
    try:
        yield factory
    finally:
        engine.dispose()


def test_global_cleanup_removes_inactive_user_history(
    session_factory: sessionmaker[Session],
) -> None:
    with session_factory() as db:
        expired = _conversation(db, "expired")
        mixed = _conversation(db, "mixed")
        fresh = _conversation(db, "fresh")
        old_time = utc_now() - timedelta(days=31)
        fresh_time = utc_now() - timedelta(days=29)
        db.add_all(
            [
                AIConversationMessage(
                    conversation_id=expired.id,
                    role="user",
                    content="expired question",
                    created_at=old_time,
                ),
                AIConversationMessage(
                    conversation_id=expired.id,
                    role="assistant",
                    content="expired answer",
                    created_at=old_time,
                ),
                AIConversationMessage(
                    conversation_id=mixed.id,
                    role="user",
                    content="old mixed question",
                    created_at=old_time,
                ),
                AIConversationMessage(
                    conversation_id=mixed.id,
                    role="assistant",
                    content="fresh mixed answer",
                    created_at=fresh_time,
                ),
                AIConversationMessage(
                    conversation_id=fresh.id,
                    role="user",
                    content="fresh question",
                    created_at=fresh_time,
                ),
            ]
        )
        db.commit()

        deleted = ConversationService(db).cleanup_expired_globally()

        assert deleted == {
            "conversation_messages": 3,
            "conversations": 1,
        }
        assert db.get(AIConversation, expired.id) is None
        assert db.get(AIConversation, mixed.id) is not None
        assert db.get(AIConversation, fresh.id) is not None
        assert db.scalar(
            select(func.count()).select_from(AIConversationMessage)
        ) == 2

    with session_factory() as verification_db:
        assert verification_db.get(AIConversation, expired.id) is None


@pytest.mark.parametrize(
    ("features_enabled", "job_expected"),
    [(False, False), (True, True)],
)
def test_scheduler_wires_one_global_retention_job(
    session_factory: sessionmaker[Session],
    monkeypatch: pytest.MonkeyPatch,
    features_enabled: bool,
    job_expected: bool,
) -> None:
    monkeypatch.setenv(
        "FEISHU_USER_FEATURES_ENABLED",
        str(features_enabled).lower(),
    )
    monkeypatch.setattr("app.core.database.SessionLocal", session_factory)
    get_settings.cache_clear()
    app = FastAPI()
    scheduler = create_scheduler(app)

    job = scheduler.get_job("personalization_retention_cleanup")

    assert (job is not None) is job_expected
    if job is None:
        return
    assert job.trigger.interval == timedelta(minutes=1)

    with session_factory() as db:
        expired = _conversation(db, "scheduled")
        db.add(
            AIConversationMessage(
                conversation_id=expired.id,
                role="user",
                content="expired scheduled message",
                created_at=utc_now() - timedelta(days=31),
            )
        )
        db.commit()
        expired_id = expired.id

    job.func()

    assert app.state.last_personalization_retention_cleanup == {
        "conversation_messages": 1,
        "conversations": 1,
    }
    with session_factory() as db:
        assert db.get(AIConversation, expired_id) is None


def _conversation(db: Session, suffix: str) -> AIConversation:
    owner = FeishuUser(
        code=f"FSU-{suffix}",
        tenant_key=f"tenant-{suffix}",
        open_id=f"open-{suffix}",
    )
    db.add(owner)
    db.flush()
    conversation = AIConversation(
        owner_id=owner.id,
        chat_type="private",
        chat_key=f"chat-{suffix}",
    )
    db.add(conversation)
    db.flush()
    return conversation