xiaobai-review/tests/test_database_migrations.py

from __future__ import annotations

import sqlite3
import tempfile
import unittest
from pathlib import Path

from backend.database import Migration, MigrationError, MigrationRunner
from database import ReviewDatabase


class DatabaseMigrationTests(unittest.TestCase):
    def test_fresh_database_records_the_adopted_schema_once(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "review.db"
            database = ReviewDatabase(path)
            with database.connect() as connection:
                rows = connection.execute(
                    "SELECT version, name FROM schema_migrations"
                ).fetchall()
                self.assertEqual(
                    [(row["version"], row["name"]) for row in rows],
                    [
                        ("0001", "adopt_legacy_schema"),
                        ("0002", "create_job_runs"),
                        ("0003", "extend_llm_audit"),
                        ("0004", "add_mentor_note"),
                        ("0005", "create_account_switch_grants"),
                    ],
                )
                columns = {
                    str(row["name"])
                    for row in connection.execute(
                        "PRAGMA table_info(mentor_preferences)"
                    )
                }
                self.assertIn("note", columns)
            ReviewDatabase(path)
            with database.connect() as connection:
                count = connection.execute(
                    "SELECT COUNT(*) AS count FROM schema_migrations"
                ).fetchone()["count"]
            self.assertEqual(count, 5)

    def test_database_with_recorded_0004_and_note_column_starts_without_reapply(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "review.db"
            database = ReviewDatabase(path)
            with database.connect() as connection:
                note_rows = [
                    str(row["name"])
                    for row in connection.execute(
                        "PRAGMA table_info(mentor_preferences)"
                    )
                ]
            self.assertIn("note", note_rows)
            ReviewDatabase(path)
            with database.connect() as connection:
                count = connection.execute(
                    "SELECT COUNT(*) AS count FROM schema_migrations"
                ).fetchone()["count"]
            self.assertEqual(count, 5)

    def test_old_database_without_0004_upgrades_and_adds_note_column(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "review.db"
            database = ReviewDatabase(path)
            with database.connect() as connection:
                connection.execute(
                    "DELETE FROM schema_migrations WHERE version = '0004'"
                )
                connection.execute(
                    "ALTER TABLE mentor_preferences DROP COLUMN note"
                )
            ReviewDatabase(path)
            with database.connect() as connection:
                versions = {
                    str(row["version"])
                    for row in connection.execute(
                        "SELECT version FROM schema_migrations"
                    )
                }
                note_rows = [
                    str(row["name"])
                    for row in connection.execute(
                        "PRAGMA table_info(mentor_preferences)"
                    )
                ]
            self.assertEqual(versions, {"0001", "0002", "0003", "0004", "0005"})
            self.assertIn("note", note_rows)

    def test_database_with_unknown_migration_is_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "review.db"
            database = ReviewDatabase(path)
            with database.connect() as connection:
                connection.execute(
                    """
                    INSERT INTO schema_migrations
                        (version, name, checksum, applied_at)
                    VALUES ('9999', 'unknown_legacy', 'x', '2026-08-01T00:00:00+00:00')
                    """
                )
            with self.assertRaises(MigrationError):
                ReviewDatabase(path)

    def test_connection_factory_enables_required_pragmas(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            database = ReviewDatabase(Path(root) / "review.db")
            with database.connect() as connection:
                self.assertEqual(connection.execute("PRAGMA foreign_keys").fetchone()[0], 1)
                self.assertEqual(connection.execute("PRAGMA journal_mode").fetchone()[0], "wal")
                self.assertEqual(connection.execute("PRAGMA busy_timeout").fetchone()[0], 20000)

    def test_failed_migration_rolls_back_and_is_not_recorded(self) -> None:
        connection = sqlite3.connect(":memory:")
        self.addCleanup(connection.close)
        connection.row_factory = sqlite3.Row

        def fail(conn: sqlite3.Connection) -> None:
            conn.execute("CREATE TABLE should_rollback (id INTEGER)")
            raise RuntimeError("stop")

        migration = Migration("9000", "failure", fail, "failure:v1")
        with self.assertRaises(MigrationError):
            MigrationRunner().apply(connection, (migration,))
        tables = {
            row["name"]
            for row in connection.execute(
                "SELECT name FROM sqlite_master WHERE type = 'table'"
            )
        }
        self.assertNotIn("should_rollback", tables)
        self.assertEqual(
            connection.execute("SELECT COUNT(*) FROM schema_migrations").fetchone()[0],
            0,
        )

    def test_applied_migration_checksum_is_immutable(self) -> None:
        connection = sqlite3.connect(":memory:")
        self.addCleanup(connection.close)
        connection.row_factory = sqlite3.Row
        first = Migration("9001", "example", lambda conn: None, "example:v1")
        changed = Migration("9001", "example", lambda conn: None, "example:v2")
        runner = MigrationRunner()
        runner.apply(connection, (first,))
        with self.assertRaises(MigrationError):
            runner.apply(connection, (changed,))


if __name__ == "__main__":
    unittest.main()