teleo-infrastructure/tests/test_verify_postgres_parity_manifest.py

from __future__ import annotations

import json
import subprocess
from pathlib import Path


def manifest_rows(
    *,
    database: str,
    row_count: int = 2,
    rowset_md5: str = "abc",
    elapsed_ms: float = 4.0,
    server_address: str | None = None,
    ssl: bool = False,
):
    return [
        {
            "kind": "identity",
            "database": database,
            "current_user": "postgres",
            "server_version_num": 160014,
            "transaction_read_only": "on",
            "server_address": server_address,
            "ssl": ssl,
        },
        {"kind": "schemas", "items": ["kb_stage", "public"]},
        {"kind": "extensions", "items": [{"name": "plpgsql", "version": "1.0"}]},
        {
            "kind": "application_roles",
            "items": [{"name": "kb_apply", "can_login": True, "superuser": False}],
        },
        {"kind": "columns", "items": [{"schema": "public", "table": "claims", "name": "id"}]},
        {"kind": "constraints", "items": [{"schema": "public", "table": "claims", "name": "claims_pkey"}]},
        {"kind": "indexes", "items": [{"schema": "public", "table": "claims", "name": "claims_pkey"}]},
        {"kind": "sequences", "items": []},
        {"kind": "views", "items": []},
        {"kind": "functions", "items": []},
        {"kind": "triggers", "items": []},
        {"kind": "types", "items": []},
        {"kind": "policies", "items": []},
        {
            "kind": "table",
            "schema": "public",
            "table": "claims",
            "row_count": row_count,
            "rowset_md5": rowset_md5,
        },
        {"kind": "performance", "query": "count_claims", "elapsed_ms": elapsed_ms, "observed_rows": row_count},
    ]


def write_jsonl(path: Path, rows) -> None:
    path.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n")


def run_verifier(tmp_path: Path, source_rows, target_rows, *, scope: str = "local_restore", connectivity=None):
    source = tmp_path / "source.jsonl"
    target = tmp_path / "target.jsonl"
    output = tmp_path / "result.json"
    write_jsonl(source, source_rows)
    write_jsonl(target, target_rows)
    command = [
        "python3",
        "ops/verify_postgres_parity_manifest.py",
        "--source",
        str(source),
        "--target",
        str(target),
        "--scope",
        scope,
        "--output",
        str(output),
    ]
    if connectivity is not None:
        connectivity_path = tmp_path / "connectivity.json"
        connectivity_path.write_text(json.dumps(connectivity))
        command.extend(["--connectivity-proof", str(connectivity_path)])
    completed = subprocess.run(command, text=True, capture_output=True, check=False)
    return completed, json.loads(output.read_text())


def test_matching_local_manifests_pass(tmp_path: Path) -> None:
    completed, payload = run_verifier(
        tmp_path,
        manifest_rows(database="teleo"),
        manifest_rows(database="teleo_copy"),
    )

    assert completed.returncode == 0, completed.stderr
    assert payload["status"] == "pass"
    assert payload["details"]["source_total_rows"] == 2
    assert payload["private_connectivity"]["status"] == "not_applicable_local_restore"


def test_row_hash_and_role_mismatches_fail(tmp_path: Path) -> None:
    target = manifest_rows(database="teleo_copy", row_count=1, rowset_md5="different")
    target[3]["items"] = []
    completed, payload = run_verifier(tmp_path, manifest_rows(database="teleo"), target)

    assert completed.returncode == 1
    assert "table_content_mismatches=1" in payload["problems"]
    assert "application_role_mismatches=1" in payload["problems"]


def test_gcp_scope_requires_private_connectivity_receipt(tmp_path: Path) -> None:
    completed, payload = run_verifier(
        tmp_path,
        manifest_rows(database="teleo"),
        manifest_rows(database="teleo_gcp_copy"),
        scope="gcp_staging",
    )

    assert completed.returncode == 1
    assert "gcp_private_connectivity_proof_missing" in payload["problems"]


def test_gcp_scope_accepts_matching_private_connectivity_receipt(tmp_path: Path) -> None:
    completed, payload = run_verifier(
        tmp_path,
        manifest_rows(database="teleo"),
        manifest_rows(database="teleo_gcp_copy", server_address="10.61.0.3/32", ssl=True),
        scope="gcp_staging",
        connectivity={
            "status": "pass",
            "private_connectivity": True,
            "public_ip_disabled": True,
            "source_compute": "teleo-staging-1",
            "target_database": "teleo_gcp_copy",
            "server_address": "10.61.0.3",
            "ssl": True,
        },
    )

    assert completed.returncode == 0, completed.stderr
    assert payload["status"] == "pass"
    assert payload["private_connectivity"]["proof"]["source_compute"] == "teleo-staging-1"


def test_gcp_scope_rejects_receipt_for_a_different_server(tmp_path: Path) -> None:
    completed, payload = run_verifier(
        tmp_path,
        manifest_rows(database="teleo"),
        manifest_rows(database="teleo_gcp_copy", server_address="10.61.0.3", ssl=True),
        scope="gcp_staging",
        connectivity={
            "status": "pass",
            "private_connectivity": True,
            "public_ip_disabled": True,
            "source_compute": "teleo-staging-1",
            "target_database": "teleo_gcp_copy",
            "server_address": "10.61.0.4",
            "ssl": True,
        },
    )

    assert completed.returncode == 1
    assert "gcp_connectivity_server_address_mismatch" in payload["problems"]


def test_gcp_scope_rejects_public_or_non_tls_target_manifest(tmp_path: Path) -> None:
    completed, payload = run_verifier(
        tmp_path,
        manifest_rows(database="teleo"),
        manifest_rows(database="teleo_gcp_copy", server_address="34.1.2.3", ssl=False),
        scope="gcp_staging",
        connectivity={
            "status": "pass",
            "private_connectivity": True,
            "public_ip_disabled": True,
            "source_compute": "teleo-staging-1",
            "target_database": "teleo_gcp_copy",
            "server_address": "34.1.2.3",
            "ssl": False,
        },
    )

    assert completed.returncode == 1
    assert "gcp_target_manifest_server_address_not_rfc1918" in payload["problems"]
    assert "gcp_target_manifest_tls_not_proven" in payload["problems"]
    assert "gcp_connectivity_tls_not_proven" in payload["problems"]