LangBot/tests/integration/persistence/test_migrations.py

"""
SQLite migration integration tests.

Tests real Alembic migration behavior using temporary SQLite databases.
Validates the migration workflow from .github/workflows/test-migrations.yml.

Run: uv run pytest tests/integration/persistence/test_migrations.py -q
"""

from __future__ import annotations

import json
import hashlib

import pytest
import sqlalchemy as sa
from alembic.config import Config
from alembic.script import ScriptDirectory
from sqlalchemy.ext.asyncio import create_async_engine

from langbot.pkg.entity.persistence import (
    agent as agent_models,
    agent_interaction as agent_interaction_models,
    agent_run as agent_run_models,
    agent_runner_state as agent_runner_state_models,
    bot as bot_models,
    metadata as metadata_models,
    monitoring as monitoring_models,
)
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import (
    _ALEMBIC_DIR,
    get_alembic_current,
    run_alembic_stamp,
    run_alembic_upgrade,
)


def _get_script_directory() -> ScriptDirectory:
    """Load the repository's Alembic revision graph."""
    cfg = Config()
    cfg.set_main_option('script_location', _ALEMBIC_DIR)
    return ScriptDirectory.from_config(cfg)


def _get_script_head() -> str:
    """Resolve the only Alembic head without hardcoding a revision."""
    return _get_script_directory().get_current_head()


pytestmark = pytest.mark.integration


class TestAlembicRevisionGraph:
    """Static release gates for the Alembic graph."""

    def test_revision_ids_fit_alembic_version_column_and_graph_has_one_head(self):
        script = _get_script_directory()
        revisions = list(script.walk_revisions())

        assert script.get_bases() == ['0001_baseline']
        assert script.get_heads() == ['0015_official_runner_ids']
        assert all(len(item.revision) <= 32 for item in revisions), {
            item.revision: len(item.revision) for item in revisions if len(item.revision) > 32
        }


@pytest.fixture
def sqlite_db_url(tmp_path):
    """Create SQLite URL with temporary database file."""
    db_file = tmp_path / 'test_migrations.db'
    return f'sqlite+aiosqlite:///{db_file}'


@pytest.fixture
async def sqlite_engine(sqlite_db_url):
    """Create async SQLite engine."""
    engine = create_async_engine(sqlite_db_url)
    yield engine
    await engine.dispose()


class TestSQLiteMigrationBaseline:
    """Tests for baseline stamp workflow."""

    @pytest.mark.asyncio
    async def test_baseline_stamp_sets_revision(self, sqlite_engine):
        """
        Stamp baseline on existing tables sets correct revision.

        Workflow:
        1. Create tables via Base.metadata.create_all
        2. Stamp with '0001_baseline'
        3. Verify current revision is '0001_baseline'
        """
        # Create all tables (simulates existing DB created by ORM)
        async with sqlite_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        # Stamp baseline
        await run_alembic_stamp(sqlite_engine, '0001_baseline')

        # Verify revision
        rev = await get_alembic_current(sqlite_engine)
        assert rev == '0001_baseline', f"Expected '0001_baseline', got {rev}"

    @pytest.mark.asyncio
    async def test_baseline_stamp_on_empty_db(self, sqlite_engine):
        """
        Stamp on empty database (no tables) still sets revision.

        This is an edge case - stamping without tables.
        """
        # Don't create tables - stamp directly
        await run_alembic_stamp(sqlite_engine, '0001_baseline')

        rev = await get_alembic_current(sqlite_engine)
        assert rev == '0001_baseline'


class TestSQLiteMigrationUpgrade:
    """Tests for upgrade to head workflow."""

    @pytest.mark.asyncio
    async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
        """
        Upgrade from baseline to head applies all migrations.

        Workflow:
        1. Create tables
        2. Stamp baseline
        3. Upgrade to head
        4. Verify current revision is head
        """
        # Create tables
        async with sqlite_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        # Stamp baseline
        await run_alembic_stamp(sqlite_engine, '0001_baseline')

        # Upgrade to head
        await run_alembic_upgrade(sqlite_engine, 'head')

        # Verify revision
        rev = await get_alembic_current(sqlite_engine)
        assert rev is not None, 'Expected a revision after upgrade'
        # Head should be the latest migration. Resolve the actual head from the
        # Alembic script directory instead of hardcoding a revision number, so
        # adding a new migration doesn't require editing this assertion.
        assert rev == _get_script_head(), f'Expected head {_get_script_head()}, got {rev}'

    @pytest.mark.asyncio
    async def test_upgrade_idempotent(self, sqlite_engine):
        """
        Running upgrade to head multiple times is idempotent.

        Workflow:
        1. Upgrade to head
        2. Get revision
        3. Upgrade to head again
        4. Verify same revision
        """
        # Create tables
        async with sqlite_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        # Stamp and upgrade
        await run_alembic_stamp(sqlite_engine, '0001_baseline')
        await run_alembic_upgrade(sqlite_engine, 'head')

        rev1 = await get_alembic_current(sqlite_engine)

        # Upgrade again - should be idempotent
        await run_alembic_upgrade(sqlite_engine, 'head')

        rev2 = await get_alembic_current(sqlite_engine)
        assert rev2 == rev1, f'Expected {rev1}, got {rev2}'

    @pytest.mark.asyncio
    async def test_upgrade_from_mcp_resource_branch_creates_agent_and_monitoring_schema(self, sqlite_engine):
        """The MCP branch can converge into the Agent branch and the current head."""
        await run_alembic_stamp(sqlite_engine, '0008_mcp_resource_prefs')
        await run_alembic_upgrade(sqlite_engine, 'head')

        def inspect_schema(sync_conn):
            inspector = sa.inspect(sync_conn)
            tables = set(inspector.get_table_names())
            monitoring_indexes = {
                index['name'] for index in inspector.get_indexes(monitoring_models.MonitoringToolCall.__tablename__)
            }
            return tables, monitoring_indexes

        async with sqlite_engine.connect() as conn:
            tables, monitoring_indexes = await conn.run_sync(inspect_schema)

        expected_agent_tables = {
            agent_models.Agent.__tablename__,
            agent_interaction_models.AgentInteraction.__tablename__,
            agent_run_models.AgentRun.__tablename__,
            agent_run_models.AgentRunEvent.__tablename__,
            agent_run_models.AgentRuntime.__tablename__,
            agent_runner_state_models.AgentRunnerState.__tablename__,
        }
        expected_monitoring_indexes = {index.name for index in monitoring_models.MonitoringToolCall.__table__.indexes}

        assert expected_agent_tables <= tables
        assert monitoring_models.MonitoringToolCall.__tablename__ in tables
        assert expected_monitoring_indexes <= monitoring_indexes
        assert await get_alembic_current(sqlite_engine) == _get_script_head()

    @pytest.mark.asyncio
    async def test_bot_admin_data_migrates_when_create_all_already_created_table(self, sqlite_engine):
        """0007 must migrate config admins even when the ORM created its table first."""
        config = {
            'admins': ['group_admin-1', 'person_user_with_underscore', 'malformed'],
            'preserved': True,
        }

        async with sqlite_engine.begin() as conn:
            await conn.run_sync(bot_models.Bot.__table__.create)
            await conn.run_sync(bot_models.BotAdmin.__table__.create)
            await conn.run_sync(metadata_models.Metadata.__table__.create)
            await conn.execute(
                sa.insert(bot_models.Bot).values(
                    uuid='bot-1',
                    name='Bot',
                    description='',
                    adapter='test',
                    adapter_config={},
                    enable=True,
                )
            )
            await conn.execute(
                sa.insert(bot_models.BotAdmin).values(
                    bot_uuid='bot-1',
                    launcher_type='group',
                    launcher_id='admin-1',
                )
            )
            await conn.execute(
                sa.insert(metadata_models.Metadata).values(
                    key='instance_config',
                    value=json.dumps(config),
                )
            )

        await run_alembic_stamp(sqlite_engine, '0006_normalize_mcp_remote_mode')
        await run_alembic_upgrade(sqlite_engine, '0007_add_bot_admins')

        async with sqlite_engine.connect() as conn:
            admin_rows = (
                await conn.execute(
                    sa.select(
                        bot_models.BotAdmin.launcher_type,
                        bot_models.BotAdmin.launcher_id,
                    ).order_by(bot_models.BotAdmin.launcher_type, bot_models.BotAdmin.launcher_id)
                )
            ).all()
            stored_config = (
                await conn.execute(
                    sa.select(metadata_models.Metadata.value).where(metadata_models.Metadata.key == 'instance_config')
                )
            ).scalar_one()

        assert admin_rows == [('group', 'admin-1'), ('person', 'user_with_underscore')]
        assert json.loads(stored_config) == {'preserved': True}

    @pytest.mark.asyncio
    async def test_pipeline_routing_rules_preserve_message_filters(self, sqlite_engine):
        """Every legacy routing rule keeps its matching semantics in event bindings."""
        routing_rules = [
            {
                'type': 'launcher_type',
                'operator': 'eq',
                'value': 'group',
                'pipeline_uuid': 'pipeline-group',
            },
            {
                'type': 'launcher_id',
                'operator': 'regex',
                'value': '^room-',
                'pipeline_uuid': 'pipeline-room',
            },
            {
                'type': 'message_content',
                'operator': 'contains',
                'value': 'urgent',
                'pipeline_uuid': 'pipeline-content',
            },
            {
                'type': 'message_has_element',
                'operator': 'eq',
                'value': 'Image',
                'pipeline_uuid': 'pipeline-image',
            },
            {
                'type': 'message_has_element',
                'operator': 'neq',
                'value': 'Voice',
                'pipeline_uuid': 'pipeline-no-voice',
            },
        ]

        async with sqlite_engine.begin() as conn:
            await conn.execute(
                sa.text(
                    'CREATE TABLE bots ('
                    'uuid VARCHAR(255) PRIMARY KEY, '
                    'use_pipeline_uuid VARCHAR(255), '
                    'pipeline_routing_rules JSON NOT NULL, '
                    'event_bindings JSON NOT NULL'
                    ')'
                )
            )
            await conn.execute(
                sa.text(
                    'INSERT INTO bots '
                    '(uuid, use_pipeline_uuid, pipeline_routing_rules, event_bindings) '
                    'VALUES (:uuid, :default_pipeline, :rules, :bindings)'
                ),
                {
                    'uuid': 'bot-routing',
                    'default_pipeline': 'pipeline-default',
                    'rules': json.dumps(routing_rules),
                    'bindings': '[]',
                },
            )

        await run_alembic_stamp(sqlite_engine, '0008_agent_product_surface')
        await run_alembic_upgrade(sqlite_engine, '0009_migrate_event_bindings')

        async with sqlite_engine.connect() as conn:
            raw_bindings = (
                await conn.execute(sa.text("SELECT event_bindings FROM bots WHERE uuid = 'bot-routing'"))
            ).scalar_one()

        bindings = json.loads(raw_bindings)
        filters_by_pipeline = {binding['target_uuid']: binding['filters'] for binding in bindings}
        assert filters_by_pipeline == {
            'pipeline-group': [{'field': 'chat_type', 'operator': 'eq', 'value': 'group'}],
            'pipeline-room': [{'field': 'chat_id', 'operator': 'regex', 'value': '^room-'}],
            'pipeline-content': [{'field': 'message_text', 'operator': 'contains', 'value': 'urgent'}],
            'pipeline-image': [{'field': 'message_element_types', 'operator': 'contains', 'value': 'Image'}],
            'pipeline-no-voice': [{'field': 'message_element_types', 'operator': 'not_contains', 'value': 'Voice'}],
            'pipeline-default': [],
        }

    @pytest.mark.asyncio
    async def test_official_runner_ids_migrate_without_losing_config_or_state(self, sqlite_engine):
        """0015 rewrites persisted official Runner identities and state scope keys."""
        old_dify = 'plugin:langbot/dify-agent/default'
        new_dify = 'plugin:langbot-team/DifyAgent/default'
        old_codex = 'plugin:langbot/codex-agent/default'
        new_codex = 'plugin:langbot-team/CodexAgent/default'
        binding_id = f'pipeline_pipeline-1_{old_dify}'
        scope_payload = {
            'version': 2,
            'scope': 'conversation',
            'runner_id': old_dify,
            'binding_identity': binding_id,
            'bot_id': 'bot-1',
            'workspace_id': None,
            'conversation_id': 'conversation-1',
            'thread_id': None,
        }
        old_scope_key = (
            'conversation:v2:'
            + hashlib.sha256(json.dumps(scope_payload, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
        )

        async with sqlite_engine.begin() as conn:
            await conn.execute(sa.text('CREATE TABLE legacy_pipelines (uuid TEXT PRIMARY KEY, config JSON NOT NULL)'))
            await conn.execute(
                sa.text('CREATE TABLE agents (uuid TEXT PRIMARY KEY, component_ref TEXT, config JSON NOT NULL)')
            )
            await conn.execute(
                sa.text(
                    'CREATE TABLE workflows ('
                    'uuid TEXT PRIMARY KEY, definition JSON NOT NULL, global_config JSON NOT NULL'
                    ')'
                )
            )
            await conn.execute(
                sa.text(
                    'CREATE TABLE agent_runner_state ('
                    'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_identity TEXT NOT NULL, '
                    'scope TEXT NOT NULL, scope_key TEXT NOT NULL, state_key TEXT NOT NULL, '
                    'bot_id TEXT, workspace_id TEXT, conversation_id TEXT, thread_id TEXT, '
                    'actor_type TEXT, actor_id TEXT, subject_type TEXT, subject_id TEXT, '
                    'UNIQUE(scope_key, state_key)'
                    ')'
                )
            )
            for table_name, columns in {
                'agent_run': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT',
                'agent_interaction': 'id INTEGER PRIMARY KEY, runner_id TEXT NOT NULL, binding_id TEXT',
                'event_log': 'id INTEGER PRIMARY KEY, runner_id TEXT',
                'transcript': 'id INTEGER PRIMARY KEY, runner_id TEXT',
            }.items():
                await conn.execute(sa.text(f'CREATE TABLE {table_name} ({columns})'))

            pipeline_config = {
                'ai': {
                    'runner': {'id': old_dify},
                    'runner_config': {
                        old_dify: {'api-key': 'preserved-secret'},
                        old_codex: {'workspace': 'K:/workspace'},
                    },
                }
            }
            await conn.execute(
                sa.text('INSERT INTO legacy_pipelines (uuid, config) VALUES (:uuid, :config)'),
                {'uuid': 'pipeline-1', 'config': json.dumps(pipeline_config)},
            )
            await conn.execute(
                sa.text('INSERT INTO agents (uuid, component_ref, config) VALUES (:uuid, :ref, :config)'),
                {'uuid': 'agent-1', 'ref': old_codex, 'config': json.dumps({'runner': {'id': old_codex}})},
            )
            await conn.execute(
                sa.text(
                    'INSERT INTO workflows (uuid, definition, global_config) '
                    'VALUES (:uuid, :definition, :global_config)'
                ),
                {
                    'uuid': 'workflow-1',
                    'definition': json.dumps({'runner_id': old_dify}),
                    'global_config': json.dumps({'runner_id': old_codex}),
                },
            )
            await conn.execute(
                sa.text(
                    'INSERT INTO agent_runner_state ('
                    'id, runner_id, binding_identity, scope, scope_key, state_key, bot_id, '
                    'workspace_id, conversation_id, thread_id, actor_type, actor_id, subject_type, subject_id'
                    ') VALUES ('
                    '1, :runner_id, :binding_identity, :scope, :scope_key, :state_key, :bot_id, '
                    'NULL, :conversation_id, NULL, NULL, NULL, NULL, NULL'
                    ')'
                ),
                {
                    'runner_id': old_dify,
                    'binding_identity': binding_id,
                    'scope': 'conversation',
                    'scope_key': old_scope_key,
                    'state_key': 'external.conversation_id',
                    'bot_id': 'bot-1',
                    'conversation_id': 'conversation-1',
                },
            )
            for table_name in ('agent_run', 'agent_interaction'):
                await conn.execute(
                    sa.text(
                        f'INSERT INTO {table_name} (id, runner_id, binding_id) VALUES (1, :runner_id, :binding_id)'
                    ),
                    {'runner_id': old_dify, 'binding_id': binding_id},
                )
            for table_name in ('event_log', 'transcript'):
                await conn.execute(
                    sa.text(f'INSERT INTO {table_name} (id, runner_id) VALUES (1, :runner_id)'),
                    {'runner_id': old_codex},
                )

        await run_alembic_stamp(sqlite_engine, '0014_interaction_delivery')
        await run_alembic_upgrade(sqlite_engine, 'head')

        async with sqlite_engine.connect() as conn:
            pipeline_raw = (
                await conn.execute(sa.text("SELECT config FROM legacy_pipelines WHERE uuid = 'pipeline-1'"))
            ).scalar_one()
            agent_row = (
                await conn.execute(sa.text("SELECT component_ref, config FROM agents WHERE uuid = 'agent-1'"))
            ).one()
            workflow_row = (
                await conn.execute(sa.text("SELECT definition, global_config FROM workflows WHERE uuid = 'workflow-1'"))
            ).one()
            state_row = (
                await conn.execute(
                    sa.text('SELECT runner_id, binding_identity, scope_key FROM agent_runner_state WHERE id = 1')
                )
            ).one()
            direct_runner_ids = {
                table_name: (
                    await conn.execute(sa.text(f'SELECT runner_id FROM {table_name} WHERE id = 1'))
                ).scalar_one()
                for table_name in ('agent_run', 'agent_interaction', 'event_log', 'transcript')
            }

        pipeline = json.loads(pipeline_raw)
        assert pipeline['ai']['runner']['id'] == new_dify
        assert pipeline['ai']['runner_config'][new_dify] == {'api-key': 'preserved-secret'}
        assert pipeline['ai']['runner_config'][new_codex] == {'workspace': 'K:/workspace'}
        assert old_dify not in pipeline['ai']['runner_config']
        assert agent_row.component_ref == new_codex
        assert json.loads(agent_row.config)['runner']['id'] == new_codex
        assert json.loads(workflow_row.definition)['runner_id'] == new_dify
        assert json.loads(workflow_row.global_config)['runner_id'] == new_codex
        assert state_row.runner_id == new_dify
        assert state_row.binding_identity == f'pipeline_pipeline-1_{new_dify}'
        assert state_row.scope_key != old_scope_key
        assert direct_runner_ids == {
            'agent_run': new_dify,
            'agent_interaction': new_dify,
            'event_log': new_codex,
            'transcript': new_codex,
        }
        assert await get_alembic_current(sqlite_engine) == '0015_official_runner_ids'


class TestSQLiteMigrationFreshDatabase:
    """Tests for fresh database workflow."""

    @pytest.mark.asyncio
    async def test_fresh_db_upgrade_from_scratch(self, tmp_path):
        """
        Fresh database (no tables) can be upgraded directly to head.

        Workflow:
        1. Create fresh engine with new DB file
        2. Create tables
        3. Upgrade to head
        4. Verify revision
        """
        # Use different DB file for fresh test
        fresh_db_file = tmp_path / 'test_migrations_fresh.db'
        fresh_url = f'sqlite+aiosqlite:///{fresh_db_file}'
        fresh_engine = create_async_engine(fresh_url)

        # Create tables on fresh DB
        async with fresh_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        # Upgrade to head directly (no baseline stamp)
        await run_alembic_upgrade(fresh_engine, 'head')

        # Verify revision
        rev = await get_alembic_current(fresh_engine)
        assert rev is not None, 'Expected a revision on fresh DB'

        await fresh_engine.dispose()

    @pytest.mark.asyncio
    async def test_fresh_db_without_create_all_behavior(self, tmp_path):
        """
        Fresh database without create_all - test actual behavior.

        This tests what happens when migrations run on truly empty DB.
        The behavior is determined by Alembic and migration scripts.

        EXPECTED: Either:
        1. Migration succeeds (if scripts handle empty DB)
        2. Migration fails with specific error (if scripts require tables)

        IMPORTANT: This test verifies the ACTUAL behavior, not accepting
        any arbitrary failure with try-except pass.
        """
        fresh_db_file = tmp_path / 'test_empty_migrations.db'
        fresh_url = f'sqlite+aiosqlite:///{fresh_db_file}'
        fresh_engine = create_async_engine(fresh_url)

        # Capture the actual behavior
        actual_result = None
        actual_error = None

        try:
            await run_alembic_upgrade(fresh_engine, 'head')
            rev = await get_alembic_current(fresh_engine)
            actual_result = rev
        except Exception as e:
            actual_error = e

        await fresh_engine.dispose()

        # Verify specific behavior - one of two outcomes is expected
        if actual_result is not None:
            # Migration succeeded - verify revision exists
            assert actual_result is not None, 'Revision should exist after successful migration'
        else:
            # Migration failed - verify the error type is known
            # Alembic typically raises specific errors for missing tables
            assert actual_error is not None, 'Error should be captured if migration failed'
            # Log the error type for documentation (don't silently pass)
            error_type = type(actual_error).__name__
            # Acceptable error types for empty DB scenarios
            acceptable_errors = [
                'OperationalError',  # SQLite table not found
                'ProgrammingError',  # SQLAlchemy errors
                'CommandError',  # Alembic command errors
            ]
            assert error_type in acceptable_errors, (
                f'Unexpected error type: {error_type}. '
                f'This may indicate a regression in migration behavior. '
                f'Error: {actual_error}'
            )


class TestSQLiteMigrationGetCurrent:
    """Tests for get_alembic_current behavior."""

    @pytest.mark.asyncio
    async def test_get_current_on_unstamped_db_returns_none(self, sqlite_engine):
        """
        get_alembic_current returns None for unstamped database.
        """
        # Create tables but don't stamp
        async with sqlite_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        # No stamp - should return None
        rev = await get_alembic_current(sqlite_engine)
        assert rev is None, f'Expected None for unstamped DB, got {rev}'

    @pytest.mark.asyncio
    async def test_get_current_after_stamp_returns_revision(self, sqlite_engine):
        """
        get_alembic_current returns correct revision after stamp.
        """
        async with sqlite_engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        await run_alembic_stamp(sqlite_engine, '0001_baseline')

        rev = await get_alembic_current(sqlite_engine)
        assert rev == '0001_baseline'