DataOps-platform/tests/runner/test_postgres_ledger.py

from pathlib import Path

from app.runner.ledger import PostgresTaskLedger


class Result:
    rowcount = 1


class Connection:
    def __init__(self):
        self.calls = []

    def execute(self, statement, parameters):
        self.calls.append((str(statement), parameters))
        return Result()


class Begin:
    def __init__(self, connection):
        self.connection = connection

    def __enter__(self):
        return self.connection

    def __exit__(self, *_args):
        return False


class Engine:
    def __init__(self):
        self.connection = Connection()

    def begin(self):
        return Begin(self.connection)


def test_postgres_ledger_claim_is_atomic_and_persists_only_bounded_fields():
    engine = Engine()
    ledger = PostgresTaskLedger(engine)
    binding = {
        "task_uid": "01900000-0000-7000-8000-000000000011",
        "dataflow_uid": "01900000-0000-7000-8000-000000000012",
        "deployment_id": "01900000-0000-7000-8000-000000000016",
        "environment": "production",
        "workflow_version": 7,
        "correlation_id": "01900000-0000-7000-8000-000000000013",
        "node_id": "write_orders",
        "node_type": "sql.execute",
        "data_source_uid": "01900000-0000-7000-8000-000000000014",
        "idempotency_key": "orders:2026-07-19",
    }

    assert ledger.claim(
        "01900000-0000-7000-8000-000000000015",
        binding,
        expires_at=2_000,
    )

    sql, parameters = engine.connection.calls[0]
    assert "ON CONFLICT DO NOTHING" in sql
    assert "encrypted_payload" not in sql
    assert "password" not in sql
    assert parameters["idempotency_key"] == "orders:2026-07-19"
    assert parameters["deployment_id"] == binding["deployment_id"]
    assert parameters["environment"] == "production"

    ledger.finish(
        "01900000-0000-7000-8000-000000000015",
        status="success",
        replay_http_status=200,
        replay_body={"result": {"rows_out": 3}},
    )
    finish_sql, finish_parameters = engine.connection.calls[1]
    assert "replay_digest" in finish_sql
    assert finish_parameters["replay_http_status"] == 200
    assert len(finish_parameters["replay_digest"]) == 64


def test_v52_migration_has_single_use_and_commit_outcome_constraints():
    root = Path(__file__).resolve().parents[2]
    source = (
        root
        / "migrations/versions/20260719_70_runner_task_ledger.py"
    ).read_text(encoding="utf-8")

    assert "runner_task_executions" in source
    assert "token_jti UUID PRIMARY KEY" in source
    assert "UNIQUE (task_uid)" in source
    assert "'not_committed', 'committed', 'unknown'" in source
    assert "encrypted_payload" not in source
    assert "connection_string" not in source