learningcircuit--local-deep-research/tests/database/test_migration_0003_indexes.py

"""
Tests for migration 0003: Add research indexes.

Tests cover:
- Index creation on research_tasks and research_history
- Full migration chain from empty database
- Idempotency (running migrations multiple times)
- Downgrade behavior
- Data preservation during migration
- Edge cases (partial tables, empty tables, in-memory databases)
"""

from alembic import command
import pytest
from sqlalchemy import create_engine, inspect, text

from local_deep_research.database.alembic_runner import (
    get_alembic_config,
    get_current_revision,
    get_head_revision,
    needs_migration,
    run_migrations,
)

# Expected indexes from migration 0003
RESEARCH_TASKS_INDEXES = {
    "ix_research_tasks_status": ["status"],
    "ix_research_tasks_created_at": ["created_at"],
    "idx_research_task_status_created": ["status", "created_at"],
    "idx_research_task_priority_status": ["priority", "status"],
}

RESEARCH_HISTORY_INDEXES = {
    "ix_research_history_mode": ["mode"],
    "ix_research_history_status": ["status"],
    "ix_research_history_created_at": ["created_at"],
    "idx_research_history_status_created": ["status", "created_at"],
    "idx_research_history_mode_status": ["mode", "status"],
}


def _get_indexes_by_name(engine, table_name):
    """Get a dict of {index_name: [column_names]} for a table."""
    insp = inspect(engine)
    if not insp.has_table(table_name):
        return {}
    return {
        idx["name"]: idx["column_names"]
        for idx in insp.get_indexes(table_name)
        if idx["name"] is not None
    }


def _run_upgrade_to(engine, revision):
    """Run migrations up to a specific revision."""
    config = get_alembic_config(engine)
    with engine.begin() as conn:
        config.attributes["connection"] = conn
        command.upgrade(config, revision)


def _run_downgrade_to(engine, revision):
    """Run downgrade to a specific revision."""
    config = get_alembic_config(engine)
    with engine.begin() as conn:
        config.attributes["connection"] = conn
        command.downgrade(config, revision)


@pytest.fixture
def fresh_engine(tmp_path):
    """Create a fresh SQLite engine (empty database, no tables)."""
    db_path = tmp_path / "fresh_0003_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    yield engine
    engine.dispose()


@pytest.fixture
def migrated_to_0002_engine(tmp_path):
    """Create a database migrated to revision 0002 (tables exist, no research indexes)."""
    db_path = tmp_path / "migrated_0002_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    _run_upgrade_to(engine, "0002")
    yield engine
    engine.dispose()


@pytest.fixture
def fully_migrated_engine(tmp_path):
    """Create a database migrated up to revision 0003 (this file's target).

    Originally this upgraded all the way to head, but later non-reversible
    migrations (0010 raises NotImplementedError on downgrade) make the
    downgrade tests below unrunnable when going through head. Since every
    test in this file is scoped to 0003 behaviour, stop the upgrade there.
    """
    db_path = tmp_path / "fully_migrated_0003_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    _run_upgrade_to(engine, "0003")
    yield engine
    engine.dispose()


class TestMigration0003UpgradeIndexes:
    """Tests that verify index creation on upgrade to 0003."""

    def test_creates_research_task_single_column_indexes(
        self, fully_migrated_engine
    ):
        """Verify ix_research_tasks_status and ix_research_tasks_created_at exist."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")

        assert "ix_research_tasks_status" in indexes
        assert indexes["ix_research_tasks_status"] == ["status"]

        assert "ix_research_tasks_created_at" in indexes
        assert indexes["ix_research_tasks_created_at"] == ["created_at"]

    def test_creates_research_task_composite_indexes(
        self, fully_migrated_engine
    ):
        """Verify composite indexes on research_tasks with correct column order."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")

        assert "idx_research_task_status_created" in indexes
        assert indexes["idx_research_task_status_created"] == [
            "status",
            "created_at",
        ]

        assert "idx_research_task_priority_status" in indexes
        assert indexes["idx_research_task_priority_status"] == [
            "priority",
            "status",
        ]

    def test_creates_research_history_single_column_indexes(
        self, fully_migrated_engine
    ):
        """Verify single-column indexes on research_history."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )

        assert "ix_research_history_mode" in indexes
        assert indexes["ix_research_history_mode"] == ["mode"]

        assert "ix_research_history_status" in indexes
        assert indexes["ix_research_history_status"] == ["status"]

        assert "ix_research_history_created_at" in indexes
        assert indexes["ix_research_history_created_at"] == ["created_at"]

    def test_creates_research_history_composite_indexes(
        self, fully_migrated_engine
    ):
        """Verify composite indexes on research_history with correct column order."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )

        assert "idx_research_history_status_created" in indexes
        assert indexes["idx_research_history_status_created"] == [
            "status",
            "created_at",
        ]

        assert "idx_research_history_mode_status" in indexes
        assert indexes["idx_research_history_mode_status"] == ["mode", "status"]

    def test_all_indexes_are_non_unique(self, fully_migrated_engine):
        """None of the 9 migration indexes should be unique."""
        insp = inspect(fully_migrated_engine)

        for table_name in ("research_tasks", "research_history"):
            for idx in insp.get_indexes(table_name):
                all_expected = {
                    **RESEARCH_TASKS_INDEXES,
                    **RESEARCH_HISTORY_INDEXES,
                }
                if idx["name"] in all_expected:
                    assert idx["unique"] == 0, (
                        f"Index {idx['name']} should not be unique"
                    )

    def test_total_index_count_research_tasks(self, fully_migrated_engine):
        """research_tasks should have exactly 4 new indexes from this migration."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")
        migration_indexes = {
            name for name in indexes if name in RESEARCH_TASKS_INDEXES
        }
        assert len(migration_indexes) == 4

    def test_total_index_count_research_history(self, fully_migrated_engine):
        """research_history should have exactly 5 new indexes from this migration."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )
        migration_indexes = {
            name for name in indexes if name in RESEARCH_HISTORY_INDEXES
        }
        assert len(migration_indexes) == 5


class TestMigration0003FromFreshDatabase:
    """Tests that verify the full migration chain on a fresh database."""

    def test_fresh_db_full_migration_creates_indexes(self, fresh_engine):
        """Running all migrations on empty DB should create all indexes."""
        run_migrations(fresh_engine)

        for table_name, expected_indexes in [
            ("research_tasks", RESEARCH_TASKS_INDEXES),
            ("research_history", RESEARCH_HISTORY_INDEXES),
        ]:
            indexes = _get_indexes_by_name(fresh_engine, table_name)
            for idx_name, idx_columns in expected_indexes.items():
                assert idx_name in indexes, (
                    f"Missing index {idx_name} on {table_name}"
                )
                assert indexes[idx_name] == idx_columns

    def test_head_revision_is_current(self):
        """get_head_revision() returns a real 4-digit revision id."""
        head = get_head_revision()
        assert head is not None and head.isdigit() and len(head) == 4

    def test_current_revision_is_head_after_migrate(self, fresh_engine):
        """After full migration, current revision should match head."""
        run_migrations(fresh_engine)
        assert get_current_revision(fresh_engine) == get_head_revision()

    def test_needs_migration_false_after_full_upgrade(self, fresh_engine):
        """After full migration, needs_migration() should return False."""
        run_migrations(fresh_engine)
        assert needs_migration(fresh_engine) is False


class TestMigration0003Idempotency:
    """Tests that verify safe re-runs of migrations."""

    def test_run_migrations_twice_no_error(self, fresh_engine):
        """Calling run_migrations() twice should not raise."""
        run_migrations(fresh_engine)
        run_migrations(fresh_engine)  # Should not raise

    def test_indexes_unchanged_after_double_migration(self, fresh_engine):
        """Indexes should be identical after running migrations twice."""
        run_migrations(fresh_engine)
        indexes_first = {
            "research_tasks": _get_indexes_by_name(
                fresh_engine, "research_tasks"
            ),
            "research_history": _get_indexes_by_name(
                fresh_engine, "research_history"
            ),
        }

        run_migrations(fresh_engine)
        indexes_second = {
            "research_tasks": _get_indexes_by_name(
                fresh_engine, "research_tasks"
            ),
            "research_history": _get_indexes_by_name(
                fresh_engine, "research_history"
            ),
        }

        assert indexes_first == indexes_second

    def test_pre_existing_indexes_not_duplicated(self, migrated_to_0002_engine):
        """Manually creating an index before migration should not cause duplicates."""
        engine = migrated_to_0002_engine

        # Manually create one of the indexes before migration
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE INDEX IF NOT EXISTS ix_research_tasks_status "
                    "ON research_tasks (status)"
                )
            )

        # Now run migration to 0003
        _run_upgrade_to(engine, "0003")

        # Verify no duplicate — only one index with that name
        indexes = _get_indexes_by_name(engine, "research_tasks")
        assert "ix_research_tasks_status" in indexes
        assert indexes["ix_research_tasks_status"] == ["status"]


class TestMigration0003Downgrade:
    """Tests for rollback behavior."""

    def test_downgrade_to_0002_removes_all_research_indexes(
        self, fully_migrated_engine
    ):
        """Downgrade from 0003 to 0002 should remove all 9 indexes."""
        _run_downgrade_to(fully_migrated_engine, "0002")

        # Fresh inspect after DDL
        all_expected = {**RESEARCH_TASKS_INDEXES, **RESEARCH_HISTORY_INDEXES}

        for table_name in ("research_tasks", "research_history"):
            indexes = _get_indexes_by_name(fully_migrated_engine, table_name)
            for idx_name in all_expected:
                assert idx_name not in indexes, (
                    f"Index {idx_name} should have been removed by downgrade"
                )

    def test_downgrade_preserves_tables(self, fully_migrated_engine):
        """Tables should still exist after downgrade (only indexes removed)."""
        _run_downgrade_to(fully_migrated_engine, "0002")

        insp = inspect(fully_migrated_engine)
        assert insp.has_table("research_tasks")
        assert insp.has_table("research_history")

    def test_downgrade_preserves_data(self, fully_migrated_engine):
        """Data inserted before downgrade should survive."""
        engine = fully_migrated_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_history (id, query, mode, status, created_at) "
                    "VALUES ('test-downgrade-1', 'test query', 'quick', 'completed', '2025-01-01')"
                )
            )

        _run_downgrade_to(engine, "0002")

        with engine.connect() as conn:
            result = conn.execute(
                text(
                    "SELECT id, query FROM research_history WHERE id = 'test-downgrade-1'"
                )
            ).fetchone()
            assert result is not None
            assert result[0] == "test-downgrade-1"
            assert result[1] == "test query"

    def test_downgrade_then_upgrade_roundtrip(self, fully_migrated_engine):
        """Downgrade to 0002 then upgrade back to 0003 should restore indexes."""
        engine = fully_migrated_engine

        _run_downgrade_to(engine, "0002")
        assert get_current_revision(engine) == "0002"

        _run_upgrade_to(engine, "0003")
        assert get_current_revision(engine) == "0003"

        # Verify indexes are back
        for table_name, expected_indexes in [
            ("research_tasks", RESEARCH_TASKS_INDEXES),
            ("research_history", RESEARCH_HISTORY_INDEXES),
        ]:
            indexes = _get_indexes_by_name(engine, table_name)
            for idx_name in expected_indexes:
                assert idx_name in indexes, (
                    f"Index {idx_name} not restored after roundtrip"
                )


class TestMigration0003DataPreservation:
    """Tests that verify migration is non-destructive."""

    def test_data_preserved_in_research_tasks(self, migrated_to_0002_engine):
        """Data in research_tasks should survive the migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Test Task', 'pending', 5, '2025-01-01 00:00:00')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            result = conn.execute(
                text("SELECT title, status, priority FROM research_tasks")
            ).fetchone()
            assert result is not None
            assert result[0] == "Test Task"
            assert result[1] == "pending"
            assert result[2] == 5

    def test_data_preserved_in_research_history(self, migrated_to_0002_engine):
        """Data in research_history should survive the migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_history (id, query, mode, status, created_at) "
                    "VALUES ('preserve-test', 'test query', 'detailed', 'completed', '2025-01-01')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            result = conn.execute(
                text(
                    "SELECT id, query, mode, status FROM research_history "
                    "WHERE id = 'preserve-test'"
                )
            ).fetchone()
            assert result is not None
            assert result[0] == "preserve-test"
            assert result[1] == "test query"
            assert result[2] == "detailed"
            assert result[3] == "completed"

    def test_data_queryable_after_index_creation(self, migrated_to_0002_engine):
        """Queries using indexed columns should work after migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Task A', 'completed', 10, '2025-01-01 00:00:00')"
                )
            )
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Task B', 'pending', 5, '2025-01-02 00:00:00')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            # Query using indexed column
            result = conn.execute(
                text(
                    "SELECT title FROM research_tasks WHERE status = 'completed'"
                )
            ).fetchall()
            assert len(result) == 1
            assert result[0][0] == "Task A"

            # Query using composite index columns
            result = conn.execute(
                text(
                    "SELECT title FROM research_tasks "
                    "WHERE priority = 10 AND status = 'completed'"
                )
            ).fetchall()
            assert len(result) == 1
            assert result[0][0] == "Task A"


class TestMigration0003EdgeCases:
    """Tests for edge cases and robustness."""

    def test_partial_tables_only_research_tasks(self, tmp_path):
        """Migration should work when only research_tasks exists."""
        db_path = tmp_path / "partial_tasks_only.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            # Create just research_tasks manually
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE research_tasks ("
                        "id INTEGER PRIMARY KEY, "
                        "title VARCHAR(500) NOT NULL, "
                        "status VARCHAR(50), "
                        "priority INTEGER DEFAULT 0, "
                        "created_at DATETIME"
                        ")"
                    )
                )
                # Create alembic_version and stamp at 0002
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            _run_upgrade_to(engine, "0003")

            indexes = _get_indexes_by_name(engine, "research_tasks")
            for idx_name in RESEARCH_TASKS_INDEXES:
                assert idx_name in indexes, f"Missing {idx_name}"

            # research_history indexes should not exist (table doesn't exist)
            insp = inspect(engine)
            assert not insp.has_table("research_history")
        finally:
            engine.dispose()

    def test_partial_tables_only_research_history(self, tmp_path):
        """Migration should work when only research_history exists."""
        db_path = tmp_path / "partial_history_only.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE research_history ("
                        "id VARCHAR(36) PRIMARY KEY, "
                        "query TEXT NOT NULL, "
                        "mode TEXT NOT NULL, "
                        "status TEXT NOT NULL, "
                        "created_at TEXT NOT NULL"
                        ")"
                    )
                )
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            _run_upgrade_to(engine, "0003")

            indexes = _get_indexes_by_name(engine, "research_history")
            for idx_name in RESEARCH_HISTORY_INDEXES:
                assert idx_name in indexes, f"Missing {idx_name}"

            insp = inspect(engine)
            assert not insp.has_table("research_tasks")
        finally:
            engine.dispose()

    def test_neither_table_exists(self, tmp_path):
        """Migration should be a no-op when neither table exists."""
        db_path = tmp_path / "neither_table.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            # Should not raise
            _run_upgrade_to(engine, "0003")
            assert get_current_revision(engine) == "0003"
        finally:
            engine.dispose()

    def test_empty_tables_get_indexes(self, migrated_to_0002_engine):
        """Tables with zero rows should still get indexes."""
        engine = migrated_to_0002_engine

        # Verify tables are empty
        with engine.connect() as conn:
            count = conn.execute(
                text("SELECT COUNT(*) FROM research_tasks")
            ).scalar()
            assert count == 0

        _run_upgrade_to(engine, "0003")

        indexes = _get_indexes_by_name(engine, "research_tasks")
        for idx_name in RESEARCH_TASKS_INDEXES:
            assert idx_name in indexes

    def test_in_memory_database(self):
        """Migration should work on in-memory SQLite database."""
        engine = create_engine("sqlite:///:memory:")

        try:
            run_migrations(engine)

            assert get_current_revision(engine) == get_head_revision()

            for table_name, expected_indexes in [
                ("research_tasks", RESEARCH_TASKS_INDEXES),
                ("research_history", RESEARCH_HISTORY_INDEXES),
            ]:
                indexes = _get_indexes_by_name(engine, table_name)
                for idx_name in expected_indexes:
                    assert idx_name in indexes, (
                        f"Missing {idx_name} on {table_name} in memory DB"
                    )
        finally:
            engine.dispose()