ModelForge/backend/tests/test_node_auth_scope_migration.py

from __future__ import annotations

import importlib.util
from pathlib import Path
from types import ModuleType

import pytest
import sqlalchemy as sa
from sqlalchemy import Engine
from sqlalchemy.exc import IntegrityError

from alembic.config import Config
from alembic.migration import MigrationContext
from alembic.operations import Operations
from alembic.script import ScriptDirectory
from modelforge_api.domain.release import TARGET_SCHEMA_REVISION
from modelforge_api.persistence.models import NodeCredential, NodeEnrollment

ROOT = Path(__file__).resolve().parents[2]
MIGRATION_PATH = (
    ROOT
    / "backend"
    / "alembic"
    / "versions"
    / "20260830_0023_node_auth_scopes.py"
)


def _migration() -> ModuleType:
    spec = importlib.util.spec_from_file_location("node_auth_scopes_0023", MIGRATION_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _scope_database(
    *, enrollment_scope: str = "node.enroll", credential_scope: str = "node.publish"
) -> Engine:
    engine = sa.create_engine("sqlite+pysqlite:///:memory:")
    metadata = sa.MetaData()
    enrollments = sa.Table(
        "node_enrollments",
        metadata,
        sa.Column("id", sa.String(36), primary_key=True),
        sa.Column("scope", sa.String(64), nullable=False),
    )
    credentials = sa.Table(
        "node_credentials",
        metadata,
        sa.Column("id", sa.String(36), primary_key=True),
        sa.Column("scope", sa.String(64), nullable=False),
    )
    metadata.create_all(engine)
    with engine.begin() as connection:
        connection.execute(enrollments.insert(), {"id": "enrollment", "scope": enrollment_scope})
        connection.execute(credentials.insert(), {"id": "credential", "scope": credential_scope})
    return engine


def _run_migration(module: ModuleType, engine: Engine, action: str) -> None:
    with engine.begin() as connection:
        module.op = Operations(MigrationContext.configure(connection))
        getattr(module, action)()


def _check_constraint_names(engine: Engine, table_name: str) -> set[str | None]:
    return {item["name"] for item in sa.inspect(engine).get_check_constraints(table_name)}


def test_0023_is_the_linear_auth_step_before_the_current_audit_head() -> None:
    config = Config(str(ROOT / "backend" / "alembic.ini"))
    config.set_main_option("script_location", str(ROOT / "backend" / "alembic"))
    config.set_main_option("path_separator", "os")
    scripts = ScriptDirectory.from_config(config)
    migration = _migration()

    assert scripts.get_heads() == ["20260830_0024"]
    assert migration.down_revision == "20260828_0022"
    assert scripts.get_revision("20260830_0024").down_revision == "20260830_0023"
    assert TARGET_SCHEMA_REVISION == "20260830_0024"


def test_current_model_metadata_carries_both_exact_scope_constraints() -> None:
    enrollment_constraints = {constraint.name for constraint in NodeEnrollment.__table__.constraints}
    credential_constraints = {constraint.name for constraint in NodeCredential.__table__.constraints}

    assert "ck_node_enrollment_scope" in enrollment_constraints
    assert "ck_node_credential_scope" in credential_constraints


def test_0023_upgrades_correct_rows_enforces_scopes_and_downgrades_cleanly() -> None:
    migration = _migration()
    engine = _scope_database()

    _run_migration(migration, engine, "upgrade")
    assert _check_constraint_names(engine, "node_enrollments") == {
        "ck_node_enrollment_scope"
    }
    assert _check_constraint_names(engine, "node_credentials") == {
        "ck_node_credential_scope"
    }

    metadata = sa.MetaData()
    metadata.reflect(engine)
    with engine.begin() as connection:
        with pytest.raises(IntegrityError):
            connection.execute(
                metadata.tables["node_enrollments"].insert(),
                {"id": "wrong-enrollment", "scope": "node.publish"},
            )
        with pytest.raises(IntegrityError):
            connection.execute(
                metadata.tables["node_credentials"].insert(),
                {"id": "wrong-credential", "scope": "node.enroll"},
            )

    _run_migration(migration, engine, "downgrade")
    assert _check_constraint_names(engine, "node_enrollments") == set()
    assert _check_constraint_names(engine, "node_credentials") == set()

    metadata = sa.MetaData()
    metadata.reflect(engine)
    with engine.begin() as connection:
        connection.execute(
            metadata.tables["node_enrollments"].insert(),
            {"id": "downgraded-enrollment", "scope": "node.publish"},
        )
        connection.execute(
            metadata.tables["node_credentials"].insert(),
            {"id": "downgraded-credential", "scope": "node.enroll"},
        )


@pytest.mark.parametrize(
    ("enrollment_scope", "credential_scope", "malformed_table"),
    [
        ("node.publish", "node.publish", "node_enrollments"),
        ("node.enroll", "node.enroll", "node_credentials"),
    ],
)
def test_0023_refuses_malformed_existing_scope_rows_before_ddl(
    enrollment_scope: str,
    credential_scope: str,
    malformed_table: str,
) -> None:
    migration = _migration()
    engine = _scope_database(
        enrollment_scope=enrollment_scope,
        credential_scope=credential_scope,
    )

    with pytest.raises(RuntimeError, match=malformed_table):
        _run_migration(migration, engine, "upgrade")

    assert _check_constraint_names(engine, "node_enrollments") == set()
    assert _check_constraint_names(engine, "node_credentials") == set()