PacketStation/tests/server/test_migrate.py

from __future__ import annotations

import contextlib
from collections.abc import Iterator
from pathlib import Path
from typing import Any

import pytest
from pshelpers import needs_pg

from packetstation import db, migrate
from packetstation.settings import Settings


class FakeResult:
    def __init__(self, rows: list[tuple[Any, ...]] | None = None):
        self.rows = rows or []

    def fetchall(self) -> list[tuple[Any, ...]]:
        return self.rows

    def fetchone(self) -> tuple[Any, ...] | None:
        return self.rows[0] if self.rows else None


class FakeConn:
    """Records statements; keeps schema_migrations in a dict; rolls back on error."""

    def __init__(self, fail_on: str | None = None):
        self.table: dict[str, str] = {}
        self.executed: list[str] = []
        self.fail_on = fail_on
        self._pending: dict[str, str] | None = None

    def execute(self, query: Any, params: Any = None) -> FakeResult:
        q = str(query)
        self.executed.append(q)
        if self.fail_on and self.fail_on in q:
            raise RuntimeError("boom")
        if q.startswith("SELECT name, sha256 FROM schema_migrations"):
            return FakeResult(sorted(self.table.items()))
        if q.startswith("INSERT INTO schema_migrations"):
            target = self._pending if self._pending is not None else self.table
            target[params[0]] = params[1]
        return FakeResult()

    @contextlib.contextmanager
    def transaction(self) -> Iterator[None]:
        self._pending = {}
        try:
            yield
        except Exception:
            self._pending = None
            raise
        self.table.update(self._pending)
        self._pending = None


def write(d: Path, name: str, sql: str) -> None:
    (d / name).write_text(sql, encoding="utf-8")


def test_applies_in_order_and_records_hashes(tmp_path: Path) -> None:
    write(tmp_path, "0002_two.sql", "-- two")
    write(tmp_path, "0001_one.sql", "-- one")
    write(tmp_path, "0010_ten.sql", "-- ten")
    conn = FakeConn()
    assert migrate.run(conn, tmp_path) == ["0001_one.sql", "0002_two.sql", "0010_ten.sql"]
    bodies = [q for q in conn.executed if q.startswith("-- ")]
    assert bodies == ["-- one", "-- two", "-- ten"]
    assert conn.table["0001_one.sql"] == migrate.sha256_file(tmp_path / "0001_one.sql")
    assert "CREATE TABLE IF NOT EXISTS schema_migrations" in conn.executed[1]
    assert migrate.run(conn, tmp_path) == []


def test_refuses_changed_file(tmp_path: Path) -> None:
    write(tmp_path, "0001_one.sql", "-- one")
    conn = FakeConn()
    migrate.run(conn, tmp_path)
    write(tmp_path, "0001_one.sql", "-- one, edited")
    write(tmp_path, "0002_two.sql", "-- two")
    with pytest.raises(migrate.MigrationError, match="changed after it was applied"):
        migrate.run(conn, tmp_path)
    assert "0002_two.sql" not in conn.table
    assert "SELECT pg_advisory_unlock(%s)" in conn.executed[-1]


def test_refuses_missing_and_out_of_order(tmp_path: Path) -> None:
    write(tmp_path, "0002_two.sql", "-- two")
    conn = FakeConn()
    migrate.run(conn, tmp_path)
    write(tmp_path, "0001_one.sql", "-- late")
    with pytest.raises(migrate.MigrationError, match="sorts before"):
        migrate.run(conn, tmp_path)
    (tmp_path / "0002_two.sql").unlink()
    with pytest.raises(migrate.MigrationError, match="missing"):
        migrate.run(conn, tmp_path)


def test_failed_migration_not_recorded(tmp_path: Path) -> None:
    write(tmp_path, "0001_one.sql", "-- one")
    write(tmp_path, "0002_bad.sql", "-- bad")
    conn = FakeConn(fail_on="-- bad")
    with pytest.raises(RuntimeError):
        migrate.run(conn, tmp_path)
    assert list(conn.table) == ["0001_one.sql"]


@pytest.mark.parametrize("name", ["1_x.sql", "0001-x.sql", "0001_X.sql"])
def test_bad_names(tmp_path: Path, name: str) -> None:
    write(tmp_path, name, "")
    with pytest.raises(migrate.MigrationError, match="bad migration file name"):
        migrate.discover(tmp_path)


def test_duplicate_numbers(tmp_path: Path) -> None:
    write(tmp_path, "0001_a.sql", "")
    write(tmp_path, "0001_b.sql", "")
    with pytest.raises(migrate.MigrationError, match="two migrations"):
        migrate.discover(tmp_path)


def test_status(tmp_path: Path) -> None:
    write(tmp_path, "0001_one.sql", "-- one")
    conn = FakeConn()
    assert migrate.status(conn, tmp_path)[0] is False
    migrate.run(conn, tmp_path)
    assert migrate.status(conn, tmp_path) == (True, "at head (0001_one.sql)")


def test_shipped_migrations_are_well_formed() -> None:
    files = migrate.discover(db.MIGRATIONS_DIR)
    assert files[0].name == "0001_base.sql"
    sql = files[0].sql()
    assert "schema_migrations (" not in sql  # the runner creates it
    revoke = sql.index("REVOKE UPDATE, DELETE, TRUNCATE ON audit_events FROM ps_app")
    assert sql.index("ALTER DEFAULT PRIVILEGES") < sql.index("CREATE TABLE audit_events") < revoke


def test_bootstrap_refused_on_prod(tmp_path: Path) -> None:
    s = Settings(_env_file=None, stack="prod", pg_superuser_url="postgresql://x@y/z")
    with pytest.raises(migrate.MigrationError, match="refused"):
        migrate.bootstrap_roles(s)
    with pytest.raises(migrate.MigrationError, match="PG_SUPERUSER_URL"):
        migrate.bootstrap_roles(Settings(_env_file=None, stack="dev"))


# ---------------------------------------------------------------- real PostgreSQL


@needs_pg
def test_pg_migrations_recorded(pg_settings: Settings) -> None:
    with db.connect_as("migrate", pg_settings, autocommit=True) as conn:
        assert migrate.run(conn, db.MIGRATIONS_DIR) == []
        rows = conn.execute("SELECT name, sha256 FROM schema_migrations ORDER BY name").fetchall()
    assert rows[0] == ("0001_base.sql", migrate.sha256_file(db.MIGRATIONS_DIR / "0001_base.sql"))


@needs_pg
def test_pg_refuses_changed_file(pg_settings: Settings, tmp_path: Path) -> None:
    d = tmp_path / "m"
    d.mkdir()
    write(d, "0001_base.sql", (db.MIGRATIONS_DIR / "0001_base.sql").read_text() + "\n-- edited\n")
    with (
        db.connect_as("migrate", pg_settings, autocommit=True) as conn,
        pytest.raises(migrate.MigrationError, match="changed"),
    ):
        migrate.run(conn, d)


@needs_pg
def test_pg_app_cannot_update_or_delete_audit(pg_settings: Settings) -> None:
    import psycopg

    from packetstation import audit

    with db.connect_as("app", pg_settings, autocommit=True) as conn:
        audit.append(conn, action="test.event", detail={"a": 1})
        for stmt in ("UPDATE audit_events SET action = 'x'", "DELETE FROM audit_events", "TRUNCATE audit_events"):
            with pytest.raises(psycopg.errors.InsufficientPrivilege):
                conn.execute(stmt)
        conn.execute("SELECT key FROM settings").fetchall()
        conn.execute("INSERT INTO settings (key, value) VALUES ('k', '1')")
    with (
        db.connect_as("migrate", pg_settings, autocommit=True) as conn,
        pytest.raises(psycopg.errors.InsufficientPrivilege),
    ):
        conn.execute("UPDATE audit_events SET action = 'x'")  # trigger refuses even the owner