X-Financial/server/tests/release_telemetry_migration_assertions.py

from __future__ import annotations

from typing import Any

import pytest
from sqlalchemy import inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import DBAPIError, IntegrityError

from app.models.agent_asset_release_telemetry import (
    AgentAssetReleaseAuditSample,
    AgentAssetReleaseLabel,
    AgentAssetReleaseObservation,
)

RELEASE_TELEMETRY_MODELS = (
    AgentAssetReleaseObservation,
    AgentAssetReleaseLabel,
    AgentAssetReleaseAuditSample,
)


def _assert_release_telemetry_head_schema(engine: Engine) -> None:
    inspector = inspect(engine)
    for model in RELEASE_TELEMETRY_MODELS:
        table = model.__table__
        live_columns = {
            str(column["name"]): bool(column["nullable"])
            for column in inspector.get_columns(table.name, schema="public")
        }
        declared_columns = {
            column.name: bool(column.nullable) for column in table.columns
        }
        assert live_columns == declared_columns

        live_constraint_names = {
            str(item["name"])
            for loader in (
                inspector.get_unique_constraints,
                inspector.get_check_constraints,
                inspector.get_foreign_keys,
            )
            for item in loader(table.name, schema="public")
            if item.get("name")
        }
        declared_constraint_names = {
            str(constraint.name)
            for constraint in table.constraints
            if constraint.name is not None
        }
        assert live_constraint_names == declared_constraint_names

        live_index_names = {
            str(item["name"])
            for item in inspector.get_indexes(table.name, schema="public")
            if not item.get("duplicates_constraint")
        }
        declared_index_names = {str(index.name) for index in table.indexes}
        assert live_index_names == declared_index_names

    label_foreign_keys = inspector.get_foreign_keys(
        "agent_asset_release_labels",
        schema="public",
    )
    assert len(label_foreign_keys) == 1
    foreign_key = label_foreign_keys[0]
    assert tuple(foreign_key["constrained_columns"]) == (
        "tenant_id",
        "observation_id",
        "asset_id",
        "release_id",
        "stage",
        "version",
    )
    assert foreign_key["referred_table"] == "agent_asset_release_observations"
    assert str(foreign_key.get("options", {}).get("ondelete", "")).upper() == "RESTRICT"

    audit_sample_foreign_keys = inspector.get_foreign_keys(
        "agent_asset_release_audit_samples",
        schema="public",
    )
    assert len(audit_sample_foreign_keys) == 1
    audit_sample_foreign_key = audit_sample_foreign_keys[0]
    assert tuple(audit_sample_foreign_key["constrained_columns"]) == (
        "tenant_id",
        "observation_id",
        "asset_id",
        "release_id",
        "stage",
        "version",
    )
    assert audit_sample_foreign_key["referred_table"] == (
        "agent_asset_release_observations"
    )
    assert (
        str(audit_sample_foreign_key.get("options", {}).get("ondelete", "")).upper()
        == "RESTRICT"
    )

    with engine.connect() as connection:
        trigger_counts = {
            table_name: int(
                connection.scalar(
                    text(
                        "SELECT COUNT(*) FROM pg_trigger trigger "
                        "JOIN pg_class relation ON relation.oid = trigger.tgrelid "
                        "WHERE relation.relname = :table_name "
                        "AND trigger.tgname = :trigger_name "
                        "AND NOT trigger.tgisinternal"
                    ),
                    {
                        "table_name": table_name,
                        "trigger_name": f"trg_{table_name}_append_only",
                    },
                )
                or 0
            )
            for table_name in (
                "agent_asset_release_observations",
                "agent_asset_release_labels",
                "agent_asset_release_audit_samples",
            )
        }
    assert trigger_counts == {
        "agent_asset_release_observations": 1,
        "agent_asset_release_labels": 1,
        "agent_asset_release_audit_samples": 1,
    }


def _assert_release_telemetry_runtime_invariants(engine: Engine) -> None:
    def rejected(
        connection: Any,
        statement: Any,
        parameters: dict[str, Any],
        error_type: type[DBAPIError] = IntegrityError,
    ) -> None:
        savepoint = connection.begin_nested()
        try:
            with pytest.raises(error_type):
                connection.execute(statement, parameters)
        finally:
            if savepoint.is_active:
                savepoint.rollback()

    observation_insert = text(
        """
        INSERT INTO agent_asset_release_observations (
            id, tenant_id, asset_id, release_id, stage, version, rule_code,
            business_stage, source_fingerprint, candidate_hit, baseline_hit,
            runtime_status, idempotency_key, payload_fingerprint
        ) VALUES (
            :id, :tenant_id, 'asset-a', 'release-a', 'canary', 'v1', 'TRAVEL-001',
            'pre_submit', :source_fingerprint, true, false, 'completed',
            :idempotency_key, :payload_fingerprint
        )
        """
    )
    label_insert = text(
        """
        INSERT INTO agent_asset_release_labels (
            id, tenant_id, observation_id, asset_id, release_id, stage, version,
            label, verification_source, source_event_fingerprint,
            actor_fingerprint, idempotency_key, payload_fingerprint
        ) VALUES (
            :id, :tenant_id, 'observation-a', 'asset-a', 'release-a', 'canary',
            'v1', 'confirmed', 'typed_risk_disposition', :event_fingerprint,
            :actor_fingerprint, :idempotency_key, :payload_fingerprint
        )
        """
    )
    audit_sample_insert = text(
        """
        INSERT INTO agent_asset_release_audit_samples (
            id, tenant_id, observation_id, asset_id, release_id, stage, version,
            stratum, sampling_probability_ppm, selection_score_ppm,
            source_reference_encrypted, idempotency_key, payload_fingerprint
        ) VALUES (
            :id, :tenant_id, 'observation-a', 'asset-a', 'release-a', 'canary',
            'v1', :stratum, :sampling_probability_ppm, :selection_score_ppm,
            :source_reference_encrypted, :idempotency_key, :payload_fingerprint
        )
        """
    )

    with engine.connect() as connection:
        transaction = connection.begin()
        try:
            observation_parameters = {
                "id": "observation-a",
                "tenant_id": "tenant-a",
                "source_fingerprint": "a" * 64,
                "idempotency_key": "observation-request-a",
                "payload_fingerprint": "b" * 64,
            }
            connection.execute(observation_insert, observation_parameters)
            rejected(
                connection,
                observation_insert,
                {
                    **observation_parameters,
                    "id": "observation-duplicate",
                    "payload_fingerprint": "c" * 64,
                },
            )
            audit_sample_parameters = {
                "id": "audit-sample-a",
                "tenant_id": "tenant-a",
                "stratum": "candidate_negative_random",
                "sampling_probability_ppm": 200_000,
                "selection_score_ppm": 12_345,
                "source_reference_encrypted": "encrypted:opaque",
                "idempotency_key": "audit-sample-request-a",
                "payload_fingerprint": "1" * 64,
            }
            connection.execute(audit_sample_insert, audit_sample_parameters)
            rejected(
                connection,
                audit_sample_insert,
                {
                    **audit_sample_parameters,
                    "id": "audit-sample-duplicate-observation",
                    "idempotency_key": "audit-sample-request-duplicate",
                    "payload_fingerprint": "2" * 64,
                },
            )
            rejected(
                connection,
                audit_sample_insert,
                {
                    **audit_sample_parameters,
                    "id": "audit-sample-cross-tenant",
                    "tenant_id": "tenant-b",
                    "idempotency_key": "audit-sample-request-cross-tenant",
                },
            )
            label_parameters = {
                "id": "label-a",
                "tenant_id": "tenant-a",
                "event_fingerprint": "d" * 64,
                "actor_fingerprint": "e" * 64,
                "idempotency_key": "label-request-a",
                "payload_fingerprint": "f" * 64,
            }
            connection.execute(label_insert, label_parameters)
            rejected(
                connection,
                label_insert,
                {
                    **label_parameters,
                    "id": "label-cross-tenant",
                    "tenant_id": "tenant-b",
                    "idempotency_key": "label-request-cross-tenant",
                },
            )
            rejected(
                connection,
                text(
                    """
                    INSERT INTO agent_asset_release_labels (
                        id, tenant_id, observation_id, asset_id, release_id, stage,
                        version, label, verification_source, source_event_fingerprint,
                        actor_fingerprint, idempotency_key, payload_fingerprint
                    ) VALUES (
                        'label-invalid-semantics', 'tenant-a', 'observation-a',
                        'asset-a', 'release-a', 'canary', 'v1', 'risk_present',
                        'release_review', :event_fingerprint, :actor_fingerprint,
                        'label-invalid-semantics', :payload_fingerprint
                    )
                    """
                ),
                {
                    "event_fingerprint": "3" * 64,
                    "actor_fingerprint": "4" * 64,
                    "payload_fingerprint": "5" * 64,
                },
            )
            connection.execute(
                text(
                    """
                    INSERT INTO agent_asset_release_labels (
                        id, tenant_id, observation_id, asset_id, release_id, stage,
                        version, label, verification_source, source_event_fingerprint,
                        actor_fingerprint, idempotency_key, payload_fingerprint
                    ) VALUES (
                        'label-blind-a', 'tenant-a', 'observation-a', 'asset-a',
                        'release-a', 'canary', 'v1', 'risk_present',
                        'blind_release_review', :event_fingerprint, :actor_fingerprint,
                        'label-blind-a', :payload_fingerprint
                    )
                    """
                ),
                {
                    "event_fingerprint": "6" * 64,
                    "actor_fingerprint": "7" * 64,
                    "payload_fingerprint": "8" * 64,
                },
            )
            rejected(
                connection,
                text(
                    "UPDATE agent_asset_release_observations "
                    "SET failure_code = 'tampered' WHERE id = 'observation-a'"
                ),
                {},
                DBAPIError,
            )
            rejected(
                connection,
                text("DELETE FROM agent_asset_release_labels WHERE id = 'label-a'"),
                {},
                DBAPIError,
            )
            rejected(
                connection,
                text(
                    "UPDATE agent_asset_release_audit_samples "
                    "SET selection_score_ppm = 999999 WHERE id = 'audit-sample-a'"
                ),
                {},
                DBAPIError,
            )
        finally:
            transaction.rollback()