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

"""Tests for migration 0010: Add chat tables.

Migration 0010 introduces the chat schema in its final clean shape:
(originally numbered 0009; renumbered to 0010 when main's
0009_default_fetch_mode_summary landed first.)

- chat_sessions with status as Enum (ChatSessionStatus)
- chat_messages with content NOT NULL, no CHECK
- chat_progress_steps as a separate table for transient research
  progress events (no longer mixed into chat_messages)
- research_history.chat_session_id (FK SET NULL) + step_count

The migration is fresh-install only; legacy 0007-shape dev DBs must
be recreated (or the chat tables dropped manually) before running.

Tests cover:
- Fresh-install path: chat tables exist, content is NOT NULL, no
  CHECK constraint, status is Enum-typed.
- Idempotency: re-running migrations on a head DB is a no-op.
- Downgrade is NotImplementedError.
"""

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

from sqlalchemy import text
from sqlalchemy.exc import IntegrityError

from local_deep_research.database.alembic_runner import (
    get_alembic_config,
    run_migrations,
    stamp_database,
)


_PARTIAL_UNIQUE_INDEX_NAME = "ux_research_history_chat_session_in_progress"


def _run_downgrade_to(engine, 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):
    db_path = tmp_path / "fresh_0010.db"
    engine = create_engine(f"sqlite:///{db_path}")
    yield engine
    engine.dispose()


@pytest.fixture
def fully_migrated_engine(tmp_path):
    db_path = tmp_path / "fully_migrated_0010.db"
    engine = create_engine(f"sqlite:///{db_path}")
    run_migrations(engine)
    yield engine
    engine.dispose()


class TestSchemaShape:
    """Chat schema invariants are in place after 0010 runs."""

    def test_chat_tables_exist(self, fully_migrated_engine):
        insp = inspect(fully_migrated_engine)
        for table in ("chat_sessions", "chat_messages", "chat_progress_steps"):
            assert insp.has_table(table), f"{table} missing"

    def test_chat_messages_content_is_not_null(self, fully_migrated_engine):
        insp = inspect(fully_migrated_engine)
        cols = {c["name"]: c for c in insp.get_columns("chat_messages")}
        assert "content" in cols
        assert cols["content"]["nullable"] is False

    def test_chat_messages_has_no_legacy_check(self, fully_migrated_engine):
        insp = inspect(fully_migrated_engine)
        checks = insp.get_check_constraints("chat_messages")
        names = {c.get("name") for c in checks}
        assert "ck_chat_message_has_content_source" not in names

    def test_chat_session_status_typed(self, fully_migrated_engine):
        """status is typed VARCHAR with default 'active'.

        Note: SQLAlchemy's Enum on SQLite stores as VARCHAR sized to
        the longest enum value but does NOT emit a DB-level CHECK
        unless `create_constraint=True` is explicitly set. The
        codebase relies on ORM-layer validation (ChatSessionStatus(value))
        for value enforcement — same pattern as ChatRole/ChatMessageType.
        """
        insp = inspect(fully_migrated_engine)
        cols = {c["name"]: c for c in insp.get_columns("chat_sessions")}
        assert "status" in cols
        # VARCHAR sized to the longest enum value ('archived' = 8 chars)
        type_str = str(cols["status"]["type"]).upper()
        assert "VARCHAR" in type_str

    def test_research_history_chat_session_id_present(
        self, fully_migrated_engine
    ):
        insp = inspect(fully_migrated_engine)
        cols = {c["name"] for c in insp.get_columns("research_history")}
        assert "chat_session_id" in cols
        assert "step_count" in cols

    def test_chat_progress_steps_unique_per_research_seq(
        self, fully_migrated_engine
    ):
        insp = inspect(fully_migrated_engine)
        uniques = insp.get_unique_constraints("chat_progress_steps")
        names = {u.get("name") for u in uniques}
        assert "uq_chat_progress_step_research_seq" in names

    def test_composite_indexes_exist_after_upgrade(self, fully_migrated_engine):
        """Composite (session_id, created_at) indexes serve the load-older
        pagination query in chat/service.py::get_session_messages. Without
        them, SQLite uses the single-column session_id index and sorts in
        memory — break-even at ~500 rows/session.
        """
        insp = inspect(fully_migrated_engine)

        msg_idx = {
            i["name"]: i["column_names"]
            for i in insp.get_indexes("chat_messages")
        }
        assert "ix_chat_messages_session_created" in msg_idx
        assert msg_idx["ix_chat_messages_session_created"] == [
            "session_id",
            "created_at",
        ]

        step_idx = {
            i["name"]: i["column_names"]
            for i in insp.get_indexes("chat_progress_steps")
        }
        assert "ix_chat_progress_steps_session_created" in step_idx
        assert step_idx["ix_chat_progress_steps_session_created"] == [
            "session_id",
            "created_at",
        ]


class TestIdempotency:
    """Re-running migrations on a head DB is a no-op."""

    def test_double_migrate_no_error(self, fresh_engine):
        run_migrations(fresh_engine)
        # Second run must not raise.
        run_migrations(fresh_engine)
        insp = inspect(fresh_engine)
        assert insp.has_table("chat_progress_steps")


class TestDowngrade:
    """Downgrade is not supported and raises NotImplementedError.

    Why: SQLite ALTER TABLE forbids dropping a column that is the
    target of a FOREIGN KEY definition, and alembic's batch_alter_table
    cannot rebuild research_history due to unnamed legacy constraints
    on that table. The project is dev-stage; recreate the DB to roll
    back. The parametrized stairway/down-leaves-no-residual tests in
    test_alembic_migrations.py exempt 0010 via NON_REVERSIBLE_REVISIONS.
    """

    def test_downgrade_raises_not_implemented(self, fully_migrated_engine):
        with pytest.raises(NotImplementedError):
            _run_downgrade_to(fully_migrated_engine, "0008")


class TestExistingDataBackfill:
    """Verify 0010 leaves pre-existing research_history rows in a sane state.

    Note: we cannot use `run_migrations(target="0008")` to land at the
    pre-0010 state because 0001 uses `Base.metadata.create_all` against the
    live `Base`, which already includes `chat_session_id` and `step_count`.
    Instead we hand-build a minimal pre-0010 `research_history` and stamp
    the DB at 0009 so 0010 forward runs the actual ADD COLUMN path.
    """

    def test_step_count_backfilled_for_existing_rows(self, tmp_path):
        engine = create_engine(f"sqlite:///{tmp_path}/m.db")

        # Hand-build pre-0010 research_history (subset of NOT NULL cols).
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE TABLE research_history ("
                    "  id TEXT PRIMARY KEY, "
                    "  query TEXT NOT NULL, "
                    "  mode TEXT NOT NULL, "
                    "  status TEXT NOT NULL, "
                    "  created_at TEXT NOT NULL"
                    ")"
                )
            )
            conn.execute(
                text(
                    "INSERT INTO research_history (id, query, mode, status, "
                    "created_at) VALUES ('r1', 'q', 'quick', 'completed', "
                    "'2026-01-01T00:00:00')"
                )
            )

        # Stamp at 0009 (main's settings-only fetch_mode migration) so
        # 0010 (chat tables) is the next forward step we exercise. 0010
        # is the migration that actually ADD COLUMNs onto our hand-built
        # research_history shape.
        stamp_database(engine, "0009")
        run_migrations(engine, target="head")

        with engine.connect() as conn:
            row = conn.execute(
                text(
                    "SELECT step_count, chat_session_id "
                    "FROM research_history WHERE id = 'r1'"
                )
            ).first()
            assert row is not None
            # Relies on 0010 using server_default="0" (SQL-side, applied by
            # SQLite ADD COLUMN at DDL time). If a future refactor switches
            # to Python-side default=0, this assertion fails for the pre-
            # existing 'r1' row — exactly the regression we want to catch.
            assert row.step_count == 0
            assert row.chat_session_id is None


# ---------------------------------------------------------------------------
# Partial unique index: at-most-one-in-progress per chat_session_id
#
# Originally lived in a separate migration 0011 with a separate test file;
# folded into 0010 to keep the chat schema landing in a single migration
# (chat is unreleased; one migration is easier to maintain than two).
# ---------------------------------------------------------------------------


@pytest.fixture
def fully_migrated_engine_with_fk(tmp_path):
    """Fully migrated SQLite engine with FK enforcement on every connection.

    Required for the partial-unique-index tests because they depend on
    SQLite enforcing the constraint at INSERT time, which only happens
    when ``PRAGMA foreign_keys = ON`` is active.
    """
    db_path = tmp_path / "0010_partial_unique_test.db"
    engine = create_engine(f"sqlite:///{db_path}")

    @event.listens_for(engine, "connect")
    def _enable_fk(dbapi_connection, _):
        dbapi_connection.execute("PRAGMA foreign_keys = ON")

    run_migrations(engine, target="head")
    yield engine
    engine.dispose()


def _seed_chat_session(conn, sid):
    conn.execute(
        text(
            "INSERT INTO chat_sessions "
            "(id, status, message_count, created_at) "
            "VALUES (:id, 'active', 0, '2026-01-01T00:00:00')"
        ),
        {"id": sid},
    )


def _insert_research(conn, *, rid, sid, status):
    conn.execute(
        text(
            "INSERT INTO research_history "
            "(id, query, mode, status, created_at, chat_session_id) "
            "VALUES (:rid, 'q', 'quick', :status, "
            "'2026-01-01T00:00:00', :sid)"
        ),
        {"rid": rid, "sid": sid, "status": status},
    )


class TestPartialUniqueInProgressIndex:
    """The partial unique index closes a SELECT-then-INSERT race in
    chat/routes.py. Verify the constraint actually fires at the DB."""

    def test_partial_unique_index_exists_after_upgrade(
        self, fully_migrated_engine_with_fk
    ):
        inspector = inspect(fully_migrated_engine_with_fk)
        indexes = {
            idx["name"]: idx
            for idx in inspector.get_indexes("research_history")
        }
        assert _PARTIAL_UNIQUE_INDEX_NAME in indexes
        idx = indexes[_PARTIAL_UNIQUE_INDEX_NAME]
        # SQLAlchemy's SQLite inspector returns 1 / 0 rather than True /
        # False for the unique flag, so compare on truthiness.
        assert bool(idx["unique"])
        assert idx["column_names"] == ["chat_session_id"]

    def test_second_in_progress_for_same_chat_session_blocked(
        self, fully_migrated_engine_with_fk
    ):
        engine = fully_migrated_engine_with_fk
        with engine.begin() as conn:
            _seed_chat_session(conn, "s1")
            _insert_research(conn, rid="r1", sid="s1", status="in_progress")

        with engine.connect() as conn:
            with pytest.raises(IntegrityError):
                with conn.begin():
                    _insert_research(
                        conn, rid="r2", sid="s1", status="in_progress"
                    )

    def test_completed_runs_for_same_chat_session_allowed(
        self, fully_migrated_engine_with_fk
    ):
        """Partial: only in_progress rows are unique; completed history
        of arbitrarily many runs per chat session must remain allowed."""
        engine = fully_migrated_engine_with_fk
        with engine.begin() as conn:
            _seed_chat_session(conn, "s1")
            _insert_research(conn, rid="r1", sid="s1", status="completed")
            _insert_research(conn, rid="r2", sid="s1", status="completed")
            _insert_research(conn, rid="r3", sid="s1", status="failed")

        with engine.connect() as conn:
            count = conn.execute(
                text(
                    "SELECT COUNT(*) FROM research_history "
                    "WHERE chat_session_id='s1'"
                )
            ).scalar()
            assert count == 3

    def test_in_progress_with_null_chat_session_id_unconstrained(
        self, fully_migrated_engine_with_fk
    ):
        """Partial: NULL chat_session_id rows must be unconstrained so
        non-chat research (news, scheduler, direct API) is unaffected."""
        engine = fully_migrated_engine_with_fk
        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_history "
                    "(id, query, mode, status, created_at) "
                    "VALUES (:rid, 'q', 'quick', 'in_progress', "
                    "'2026-01-01T00:00:00')"
                ),
                {"rid": "r1"},
            )
            conn.execute(
                text(
                    "INSERT INTO research_history "
                    "(id, query, mode, status, created_at) "
                    "VALUES (:rid, 'q', 'quick', 'in_progress', "
                    "'2026-01-01T00:00:00')"
                ),
                {"rid": "r2"},
            )

        with engine.connect() as conn:
            count = conn.execute(
                text(
                    "SELECT COUNT(*) FROM research_history "
                    "WHERE chat_session_id IS NULL"
                )
            ).scalar()
            assert count == 2

    def test_completing_a_run_releases_the_in_progress_slot(
        self, fully_migrated_engine_with_fk
    ):
        """After r1 transitions away from in_progress, r2 must be able
        to claim the slot."""
        engine = fully_migrated_engine_with_fk
        with engine.begin() as conn:
            _seed_chat_session(conn, "s1")
            _insert_research(conn, rid="r1", sid="s1", status="in_progress")

        with engine.begin() as conn:
            conn.execute(
                text(
                    "UPDATE research_history "
                    "SET status='completed' WHERE id='r1'"
                )
            )

        with engine.begin() as conn:
            _insert_research(conn, rid="r2", sid="s1", status="in_progress")

        with engine.connect() as conn:
            count = conn.execute(
                text(
                    "SELECT COUNT(*) FROM research_history "
                    "WHERE chat_session_id='s1' AND status='in_progress'"
                )
            ).scalar()
            assert count == 1