DataOps-platform/tests/test_data_rule_schema.py

from __future__ import annotations

import importlib.util
from pathlib import Path

import pytest

ROOT = Path(__file__).resolve().parents[1]
MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_110_ai_data_rules.py"
)
RUNTIME_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_120_rule_execution_runtime.py"
)
PLAN_STATUS_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_130_bound_plan_lifecycle.py"
)
ARTIFACT_CATALOG_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_140_rule_run_artifacts.py"
)
ARTIFACT_HANDOFF_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_150_rule_artifact_handoff_state.py"
)
RULE_EVIDENCE_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_160_rule_execution_evidence.py"
)
RULE_ATTEMPT_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_170_rule_execution_attempts.py"
)
RULE_RECONCILIATION_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_180_rule_reconciliation.py"
)
RULE_CLEANUP_LEASE_MIGRATION = (
    ROOT
    / "migrations"
    / "versions"
    / "20260723_190_rule_cleanup_leases.py"
)

EXPECTED_TABLES = {
    "data_rules",
    "data_rule_versions",
    "rule_generation_runs",
    "data_standards",
    "data_standard_versions",
    "standard_rule_bindings",
    "dataflow_versions",
    "dataflow_component_bindings",
    "rule_execution_plans",
    "rule_artifacts",
    "dataflow_deployments",
    "rule_runs",
    "rule_violation_samples",
}


def test_ai_data_rule_migration_defines_all_three_domain_layers():
    source = MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_110"' in source
    assert 'down_revision = "20260722_110"' in source
    for table in EXPECTED_TABLES:
        assert f"CREATE TABLE public.{table}" in source


def test_ai_data_rule_migration_enforces_immutable_versions_and_fixed_bindings():
    source = MIGRATION.read_text(encoding="utf-8")

    for expected in (
        "UNIQUE (rule_uid, version_no)",
        "UNIQUE (standard_uid, version_no)",
        "UNIQUE (dataflow_uid, version_no)",
        "standard_version_id UUID NOT NULL",
        "rule_version_id UUID NOT NULL",
        "dataflow_version_id UUID NOT NULL",
        "component_kind VARCHAR(30) NOT NULL",
        "package_hash CHAR(64)",
        "plan_hash CHAR(64)",
        "artifact_digest CHAR(64)",
        "context_hash CHAR(64)",
        "candidate_hash CHAR(64)",
        "workflow_version_id UUID",
    ):
        assert expected in source


def test_ai_data_rule_migration_is_forward_preserving():
    source = MIGRATION.read_text(encoding="utf-8")
    downgrade = source.split("def downgrade()", 1)[1]

    assert "DROP TABLE" not in downgrade.upper()
    assert "pass" in downgrade


def test_dataflow_draft_reservation_migration_is_single_use_and_forward_only():
    source = (
        ROOT
        / "migrations/versions/20260723_210_dataflow_draft_reservations.py"
    ).read_text(encoding="utf-8")

    assert 'revision = "20260723_210"' in source
    assert 'down_revision = "20260723_200"' in source
    assert "dataflow_draft_reservations" in source
    assert "dataflow_uid UUID NOT NULL UNIQUE" in source
    assert "actor_uid UUID NOT NULL" in source
    assert "nonce_hash CHAR(64) NOT NULL UNIQUE" in source
    assert "consumed_at TIMESTAMPTZ" in source
    assert "expires_at TIMESTAMPTZ NOT NULL" in source
    assert "forward-only" in source


def test_rule_execution_runtime_migration_adds_pinned_runtime_evidence():
    source = RUNTIME_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_120"' in source
    assert 'down_revision = "20260723_110"' in source
    for table in (
        "data_schema_snapshots",
        "dataflow_dataset_bindings",
        "rule_compile_evidence",
        "rule_test_evidence",
    ):
        assert f"CREATE TABLE public.{table}" in source
    for expected in (
        "rule_execution_plan_id UUID NOT NULL",
        "created_by UUID REFERENCES public.users(id)",
        "candidate JSONB",
        "UNIQUE (schema_ref, schema_hash)",
        "UNIQUE (dataflow_deployment_id, logical_ref)",
    ):
        assert expected in source


def test_rule_execution_runtime_migration_remains_forward_preserving():
    source = RUNTIME_MIGRATION.read_text(encoding="utf-8")
    downgrade = source.split("def downgrade()", 1)[1]

    assert "DROP TABLE" not in downgrade.upper()
    assert "pass" in downgrade


def test_bound_plan_status_migration_is_forward_only_and_keeps_compiled_valid():
    source = PLAN_STATUS_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_130"' in source
    assert 'down_revision = "20260723_120"' in source
    assert "'compiled','published','revoked'" in source
    assert "'tested'" not in source
    assert "raise RuntimeError" in source.split("def downgrade()", 1)[1]
    spec = importlib.util.spec_from_file_location(
        "bound_plan_status_migration",
        PLAN_STATUS_MIGRATION,
    )
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    with pytest.raises(RuntimeError, match="forward-only|cannot downgrade"):
        module.downgrade()


def test_rule_run_artifact_catalog_is_correlation_scoped_and_forward_preserving():
    source = ARTIFACT_CATALOG_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_140"' in source
    assert 'down_revision = "20260723_130"' in source
    assert "CREATE TABLE public.rule_run_artifacts" in source
    for expected in (
        "correlation_id UUID NOT NULL",
        "binding_id UUID NOT NULL",
        "artifact_ref VARCHAR(1000) NOT NULL",
        "artifact_digest CHAR(64) NOT NULL",
        "schema_fields JSONB NOT NULL",
        "expires_at TIMESTAMPTZ NOT NULL",
        "UNIQUE (correlation_id, binding_id, artifact_digest)",
    ):
        assert expected in source
    downgrade = source.split("def downgrade()", 1)[1]
    assert "DROP TABLE" not in downgrade.upper()
    assert "pass" in downgrade


def test_artifact_handoff_state_migration_upgrades_old_140_forward_only():
    source = ARTIFACT_HANDOFF_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_150"' in source
    assert 'down_revision = "20260723_140"' in source
    for expected in (
        "binding_hash CHAR(64)",
        "handoff_status VARCHAR(20)",
        "ready_at TIMESTAMPTZ",
        "failed_at TIMESTAMPTZ",
        "updated_at TIMESTAMPTZ",
        "UNIQUE (correlation_id, binding_id, artifact_kind)",
        "'pending','ready','failed'",
        "SELECT binding_hash",
    ):
        assert expected in source
    assert "artifact_digest" in source
    assert "RAISE EXCEPTION" in source
    downgrade = source.split("def downgrade()", 1)[1]
    assert "DROP TABLE" not in downgrade.upper()
    assert "raise RuntimeError" in downgrade
    spec = importlib.util.spec_from_file_location(
        "rule_artifact_handoff_state_migration",
        ARTIFACT_HANDOFF_MIGRATION,
    )
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    with pytest.raises(RuntimeError, match="forward-only|cannot downgrade"):
        module.downgrade()


def test_rule_evidence_upgrade_preflights_legacy_duplicates_and_oversize():
    source = RULE_EVIDENCE_MIGRATION.read_text(encoding="utf-8")

    duplicate_check = source.index("HAVING COUNT(*) > 1")
    unique_constraint = source.index(
        "ADD CONSTRAINT rule_violation_samples_rule_run_key"
    )
    assert duplicate_check < unique_constraint
    assert "duplicate rule_run_id" in source
    assert "sample_count > 100" in source
    assert "rows above 100" in source


def test_rule_attempt_migration_adds_exact_identity_replay_lease_and_receipts():
    source = RULE_ATTEMPT_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_170"' in source
    assert 'down_revision = "20260723_160"' in source
    for expected in (
        "lease_owner UUID",
        "lease_expires_at TIMESTAMPTZ",
        "evidence_digest CHAR(64)",
        "deployment_id UUID",
        "environment VARCHAR(20)",
        "replay_body JSONB",
        "replay_digest CHAR(64)",
        "CREATE TABLE public.rule_sql_staging_receipts",
        "producer_rule_run_id UUID NOT NULL",
        "relation_digest CHAR(64) NOT NULL",
        "UNIQUE (producer_rule_run_id, output_binding_id)",
    ):
        assert expected in source
    assert "DROP TABLE" not in source.split("def downgrade()", 1)[1].upper()


def test_rule_reconciliation_migration_extends_170_without_rewriting_it():
    source = RULE_RECONCILIATION_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_180"' in source
    assert 'down_revision = "20260723_170"' in source
    for expected in (
        "runner_task_executions",
        "lease_expires_at",
        "'unknown','expired'",
        "cleanup_claim UUID",
        "schema_fields JSONB",
        "rule_sql_staging_receipts",
    ):
        assert expected in source
    assert "forward-only" in source


def test_rule_cleanup_lease_migration_extends_180_forward_only():
    source = RULE_CLEANUP_LEASE_MIGRATION.read_text(encoding="utf-8")

    assert 'revision = "20260723_190"' in source
    assert 'down_revision = "20260723_180"' in source
    assert source.count("cleanup_claim_expires_at TIMESTAMPTZ") == 2
    assert "rule_violation_samples" in source
    assert "rule_sql_staging_receipts" in source
    assert source.count(
        "SET cleanup_claim_expires_at = CURRENT_TIMESTAMP"
    ) == 2
    assert source.count("WHERE cleanup_claim IS NOT NULL") == 2
    assert source.count("AND cleanup_claim_expires_at IS NULL") == 2
    assert "forward-only" in source