"""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