company-ai-platform/tests/test_feishu_personalization_migration.py

from datetime import datetime
from pathlib import Path

import pytest
from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text

from app.core.config import get_settings


def test_personalization_migration_preserves_and_classifies_legacy_data(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    database_path = tmp_path / "personalization-migration.db"
    database_url = f"sqlite:///{database_path.as_posix()}"
    monkeypatch.setenv("DATABASE_URL", database_url)
    get_settings.cache_clear()
    config = Config("alembic.ini")
    engine = create_engine(database_url)
    timestamp = datetime(2026, 7, 26, 8, 0)

    try:
        command.upgrade(config, "202607260002")
        with engine.begin() as connection:
            connection.execute(
                text(
                    """
                    INSERT INTO ai_memory_entries (
                        code,
                        fingerprint,
                        scope,
                        subject,
                        content,
                        source,
                        importance,
                        status,
                        actor,
                        created_at,
                        updated_at
                    )
                    VALUES (
                        :code,
                        :fingerprint,
                        'global',
                        :subject,
                        :content,
                        :source,
                        1,
                        'active',
                        'legacy-user',
                        :created_at,
                        :updated_at
                    )
                    """
                ),
                [
                    {
                        "code": "legacy-rule",
                        "fingerprint": "legacy-rule-fingerprint",
                        "subject": "rule",
                        "content": "Use concise Chinese.",
                        "source": "user_rule",
                        "created_at": timestamp,
                        "updated_at": timestamp,
                    },
                    {
                        "code": "legacy-auto-memory",
                        "fingerprint": "legacy-auto-fingerprint",
                        "subject": "preference",
                        "content": "Prefers market summaries.",
                        "source": "auto",
                        "created_at": timestamp,
                        "updated_at": timestamp,
                    },
                    {
                        "code": "legacy-hermes-memory",
                        "fingerprint": "legacy-hermes-fingerprint",
                        "subject": "preference",
                        "content": "Prefers short answers.",
                        "source": "hermes",
                        "created_at": timestamp,
                        "updated_at": timestamp,
                    },
                ],
            )
            connection.execute(
                text(
                    """
                    INSERT INTO market_watchlists (
                        actor,
                        symbol,
                        enabled,
                        created_at,
                        updated_at
                    )
                    VALUES (
                        'legacy-user',
                        '600000.SH',
                        1,
                        :created_at,
                        :updated_at
                    )
                    """
                ),
                {"created_at": timestamp, "updated_at": timestamp},
            )

        command.upgrade(config, "head")

        inspector = inspect(engine)
        expected_tables = {
            "feishu_app_tickets",
            "feishu_admin_bootstrap_tombstones",
            "feishu_users",
            "user_preferences",
            "ai_conversations",
            "ai_conversation_messages",
            "personal_data_erasure_requests",
            "push_subscriptions",
            "push_deliveries",
        }
        assert expected_tables <= set(inspector.get_table_names())
        tombstone_columns = {
            column["name"]
            for column in inspector.get_columns(
                "feishu_admin_bootstrap_tombstones"
            )
        }
        assert tombstone_columns == {"id", "identity_hash", "created_at"}

        memory_columns = {
            column["name"] for column in inspector.get_columns("ai_memory_entries")
        }
        assert {"owner_id", "kind"} <= memory_columns
        memory_indexes = {
            index["name"]: index for index in inspector.get_indexes("ai_memory_entries")
        }
        assert memory_indexes["ix_ai_memory_entries_fingerprint"]["unique"] == 0
        assert {
            "ix_ai_memory_entries_kind",
            "ix_ai_memory_entries_owner_id",
        } <= set(memory_indexes)
        memory_constraints = {
            constraint["name"]
            for constraint in inspector.get_unique_constraints("ai_memory_entries")
        }
        assert "uq_ai_memory_owner_fingerprint" in memory_constraints
        memory_foreign_keys = {
            foreign_key["name"]
            for foreign_key in inspector.get_foreign_keys("ai_memory_entries")
        }
        assert "fk_ai_memory_entries_owner_id" in memory_foreign_keys

        with engine.connect() as connection:
            memories = {
                row.code: row
                for row in connection.execute(
                    text(
                        """
                        SELECT code, owner_id, kind, source, status
                        FROM ai_memory_entries
                        ORDER BY code
                        """
                    )
                )
            }
            watchlist_owner_id = connection.scalar(
                text(
                    """
                    SELECT owner_id
                    FROM market_watchlists
                    WHERE actor = 'legacy-user' AND symbol = '600000.SH'
                    """
                )
            )

        legacy_rule = memories["legacy-rule"]
        assert legacy_rule.owner_id is None
        assert legacy_rule.kind == "company_rule"
        assert legacy_rule.source == "legacy_company"
        assert legacy_rule.status == "active"
        for code in ("legacy-auto-memory", "legacy-hermes-memory"):
            assert memories[code].owner_id is None
            assert memories[code].kind == "memory"
            assert memories[code].status == "archived"
        assert watchlist_owner_id is None

        command.check(config)
    finally:
        engine.dispose()
        get_settings.cache_clear()