DocsGPT/tests/storage/db/test_migration_0011.py

"""Migration round-trip test for 0011_schedules_nullable_agent."""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

import pytest
from sqlalchemy import text


pytestmark = pytest.mark.integration


def _alembic_ini() -> Path:
    return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"


def _run_alembic(url: str, *args: str) -> None:
    """Run ``alembic`` against ``url``."""
    subprocess.check_call(
        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
        timeout=60,
        env={**os.environ, "POSTGRES_URI": url},
    )


def _column_is_nullable(conn, table: str, column: str) -> bool:
    row = conn.execute(
        text(
            "SELECT is_nullable FROM information_schema.columns "
            "WHERE table_name = :t AND column_name = :c"
        ),
        {"t": table, "c": column},
    ).fetchone()
    return row is not None and row[0] == "YES"


def _alembic_version(conn) -> str:
    return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()


_0011_REVISION = "0011_schedules_nullable_agent"
_0010_REVISION = "0010_schedules"


class TestMigration0011RoundTrip:
    def test_head_state_has_nullable_agent_id(self, pg_engine):
        # pg_engine upgrades to ``head`` which includes 0011.
        with pg_engine.connect() as conn:
            assert _alembic_version(conn) >= _0011_REVISION
            assert _column_is_nullable(conn, "schedules", "agent_id")
            assert _column_is_nullable(conn, "schedule_runs", "agent_id")

    def test_agentless_insert_succeeds_at_head(self, pg_engine):
        # Sanity: NULL agent_id is accepted by the schema.
        with pg_engine.begin() as conn:
            row = conn.execute(
                text(
                    """
                    INSERT INTO schedules
                        (user_id, agent_id, trigger_type, run_at, instruction)
                    VALUES (
                        'u-agentless', NULL, 'once',
                        now() + interval '1 hour', 'go'
                    )
                    RETURNING id
                    """
                ),
            ).fetchone()
            schedule_id = str(row[0])
            conn.execute(
                text(
                    """
                    INSERT INTO schedule_runs
                        (schedule_id, user_id, agent_id, scheduled_for)
                    VALUES (CAST(:s AS uuid), 'u-agentless', NULL, now())
                    """
                ),
                {"s": schedule_id},
            )
            cnt = conn.execute(
                text(
                    "SELECT count(*) FROM schedules "
                    "WHERE agent_id IS NULL AND user_id = 'u-agentless'"
                )
            ).scalar()
            assert cnt == 1

    def test_downgrade_restores_not_null_when_no_nulls(self, pg_engine):
        url = pg_engine.url.render_as_string(hide_password=False)
        # Downgrade to 0010: agent_id NOT NULL is restored.
        _run_alembic(url, "downgrade", _0010_REVISION)
        with pg_engine.connect() as conn:
            assert _alembic_version(conn) == _0010_REVISION
            assert not _column_is_nullable(conn, "schedules", "agent_id")
            assert not _column_is_nullable(conn, "schedule_runs", "agent_id")
        _run_alembic(url, "upgrade", "head")

    def test_downgrade_fails_when_agentless_rows_present(self, pg_engine):
        """Downgrade with NULL agent_id rows raises loudly (no data loss)."""
        # Start at head (already there). Insert an agentless row.
        with pg_engine.begin() as conn:
            conn.execute(
                text(
                    """
                    INSERT INTO schedules
                        (user_id, agent_id, trigger_type, run_at, instruction)
                    VALUES (
                        'u-block', NULL, 'once',
                        now() + interval '1 hour', 'go'
                    )
                    """
                )
            )
        url = pg_engine.url.render_as_string(hide_password=False)
        with pytest.raises(subprocess.CalledProcessError):
            _run_alembic(url, "downgrade", _0010_REVISION)
        # Clean up so downstream tests aren't affected.
        with pg_engine.begin() as conn:
            conn.execute(
                text("DELETE FROM schedules WHERE user_id = 'u-block'")
            )

    def test_agentless_rows_survive_unrelated_agent_delete(self, pg_engine):
        """Agentless rows (agent_id IS NULL) aren't matched by FK CASCADE."""
        with pg_engine.begin() as conn:
            agent_id = conn.execute(
                text(
                    "INSERT INTO agents (user_id, name, status) "
                    "VALUES ('u1', 'a', 'draft') RETURNING id"
                )
            ).fetchone()[0]
            conn.execute(
                text(
                    "INSERT INTO schedules "
                    "(user_id, agent_id, trigger_type, run_at, instruction) "
                    "VALUES ('u-agentless', NULL, 'once', "
                    "now() + interval '1 hour', 'survive')"
                )
            )
            conn.execute(
                text("DELETE FROM agents WHERE id = CAST(:a AS uuid)"),
                {"a": str(agent_id)},
            )
            count = conn.execute(
                text(
                    "SELECT count(*) FROM schedules "
                    "WHERE user_id = 'u-agentless'"
                )
            ).scalar()
            assert count == 1

    def test_full_round_trip_lands_back_at_head(self, pg_engine):
        url = pg_engine.url.render_as_string(hide_password=False)
        _run_alembic(url, "downgrade", _0010_REVISION)
        _run_alembic(url, "upgrade", _0011_REVISION)
        with pg_engine.connect() as conn:
            assert _alembic_version(conn) == _0011_REVISION
            assert _column_is_nullable(conn, "schedules", "agent_id")
            assert _column_is_nullable(conn, "schedule_runs", "agent_id")
        _run_alembic(url, "upgrade", "head")