mimic/backend/tests/test_migration_0007_c2.py

"""Migration round-trip test for 0007_c2_task_mapping_applied.

Verifies that upgrade() adds the mapping_applied column and downgrade() removes it.
Uses the resolved-path pattern per lessons.md Sprint 4.
"""
from __future__ import annotations

import importlib.util
from pathlib import Path

from alembic.operations import Operations
from alembic.runtime.migration import MigrationContext
from sqlalchemy import create_engine, inspect, text


def _load_migration(name: str):
    versions_dir = Path(__file__).resolve().parent.parent / "migrations" / "versions"
    path = versions_dir / name
    spec = importlib.util.spec_from_file_location(name.removesuffix(".py"), path)
    assert spec and spec.loader
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)  # type: ignore[union-attr]
    return mod


def _fresh_engine_with_c2_task():
    """In-memory SQLite with c2_task already created (as left by 0006 upgrade)."""
    engine = create_engine("sqlite:///:memory:")
    with engine.begin() as conn:
        conn.execute(text("""
            CREATE TABLE c2_task (
                id INTEGER PRIMARY KEY,
                simulation_id INTEGER NOT NULL,
                mythic_task_display_id INTEGER NOT NULL,
                callback_display_id INTEGER NOT NULL,
                command TEXT NOT NULL,
                params TEXT,
                status TEXT NOT NULL,
                completed BOOLEAN NOT NULL DEFAULT 0,
                output TEXT,
                source TEXT NOT NULL,
                created_at DATETIME NOT NULL,
                completed_at DATETIME
            )
        """))
    return engine


def _run_upgrade(engine, migration_mod):
    with engine.begin() as conn:
        ctx = MigrationContext.configure(conn)
        ops = Operations(ctx)
        ops._install_proxy()  # type: ignore[attr-defined]
        try:
            migration_mod.upgrade()
        finally:
            ops._remove_proxy()  # type: ignore[attr-defined]


def _run_downgrade(engine, migration_mod):
    with engine.begin() as conn:
        ctx = MigrationContext.configure(conn)
        ops = Operations(ctx)
        ops._install_proxy()  # type: ignore[attr-defined]
        try:
            migration_mod.downgrade()
        finally:
            ops._remove_proxy()  # type: ignore[attr-defined]


class TestMigration0007Upgrade:
    def test_mapping_applied_column_added(self):
        engine = _fresh_engine_with_c2_task()
        mod = _load_migration("0007_c2_task_mapping_applied.py")
        _run_upgrade(engine, mod)

        insp = inspect(engine)
        cols = {c["name"] for c in insp.get_columns("c2_task")}
        assert "mapping_applied" in cols

    def test_mapping_applied_defaults_to_false(self):
        engine = _fresh_engine_with_c2_task()
        mod = _load_migration("0007_c2_task_mapping_applied.py")

        # Insert a row before upgrading (no mapping_applied column yet).
        with engine.begin() as conn:
            conn.execute(text(
                "INSERT INTO c2_task "
                "(simulation_id, mythic_task_display_id, callback_display_id, "
                "command, status, completed, source, created_at) "
                "VALUES (1, 1000, 1, 'whoami', 'submitted', 0, 'mimic', '2026-01-01')"
            ))

        _run_upgrade(engine, mod)

        with engine.begin() as conn:
            row = conn.execute(
                text("SELECT mapping_applied FROM c2_task WHERE id = 1")
            ).fetchone()
        assert row is not None
        # SQLite stores booleans as 0/1.
        assert row[0] == 0 or row[0] is False


class TestMigration0007Downgrade:
    def test_downgrade_removes_mapping_applied(self):
        engine = _fresh_engine_with_c2_task()
        mod = _load_migration("0007_c2_task_mapping_applied.py")
        _run_upgrade(engine, mod)
        _run_downgrade(engine, mod)

        insp = inspect(engine)
        cols = {c["name"] for c in insp.get_columns("c2_task")}
        assert "mapping_applied" not in cols

    def test_downgrade_does_not_drop_other_columns(self):
        engine = _fresh_engine_with_c2_task()
        mod = _load_migration("0007_c2_task_mapping_applied.py")
        _run_upgrade(engine, mod)
        _run_downgrade(engine, mod)

        insp = inspect(engine)
        cols = {c["name"] for c in insp.get_columns("c2_task")}
        assert {"id", "simulation_id", "command", "status", "completed"} <= cols