from __future__ import annotations import hashlib import unittest from datetime import date, datetime, timedelta, timezone from typing import Any from arr_ingestion.contracts import IngestionError from arr_ingestion.direct_contracts import ( DIRECT_CONTRACT_VERSION, DirectSubmissionReceipt, DirectSubmissionRequest, ) from arr_ingestion.direct_postgres import PostgresDirectIngestionRepository from arr_ingestion.postgres import DatabaseConfig TOKEN = "T" * 43 JOB_ID = "arrjob-direct-postgres" RULE_HASH = "a" * 64 SOURCE_HASH = "b" * 64 SOURCE_KEY = ( "arr/jobs/arrjob-direct-postgres/attempts/0001/committed/" "source_xml/source.xml" ) def minimal_payload(*, marker: str | None = None) -> dict[str, Any]: payload: dict[str, Any] = { "status": "success", "activation_eligible": True, "business_date": "2026-07-27", "processor_version": "3.0.0", "rule_set_sha256": RULE_HASH, "result_schema_version": "3.0", "records": [], } if marker is not None: payload["marker"] = marker return payload def request(*, marker: str | None = None) -> DirectSubmissionRequest: return DirectSubmissionRequest.from_dict( { "submission_grant": TOKEN, "job_id": JOB_ID, "attempt_no": 1, "payload": minimal_payload(marker=marker), } ) def receipt(value: DirectSubmissionRequest) -> dict[str, Any]: return DirectSubmissionReceipt( "committed", JOB_ID, 1, date(2026, 7, 27), 71, 2, value.record_count, ).to_dict() class ScriptedCursor: def __init__(self, responses: list[tuple[str, Any]]): self.responses = responses self.last_query = "" self.executed: list[tuple[str, tuple[Any, ...] | None]] = [] def execute( self, query: str, params: tuple[Any, ...] | None = None, ) -> None: normalized = " ".join(query.split()) if params is not None and normalized.count("%s") != len(params): raise AssertionError( f"placeholder mismatch: {normalized.count('%s')} != {len(params)}" ) self.last_query = normalized self.executed.append((normalized, params)) def fetchone(self): for index, (marker, value) in enumerate(self.responses): if marker in self.last_query: self.responses.pop(index) return value return None class DirectPostgresTests(unittest.TestCase): def test_readiness_requires_every_direct_and_finance_relation(self) -> None: cursor = ScriptedCursor( [ ( "to_regclass('ingestion.result_submission_grants')", ("grants", "submissions", "versions", "records", "current"), ) ] ) PostgresDirectIngestionRepository._assert_ready(cursor) missing = ScriptedCursor( [ ( "to_regclass('ingestion.result_submission_grants')", ("grants", None, "versions", "records", "current"), ) ] ) with self.assertRaises(IngestionError) as raised: PostgresDirectIngestionRepository._assert_ready(missing) self.assertEqual(raised.exception.code, "DATABASE_MIGRATION_MISSING") def test_grant_persists_only_hash_and_switches_attempt_to_direct_mode(self) -> None: expires_at = datetime.now(timezone.utc) + timedelta(minutes=10) cursor = ScriptedCursor( [ ( "FROM ingestion.processing_runs AS run", (11, "opera_daily", "running", "artifact_callback", 12, "running"), ), ("FROM ingestion.processing_deliveries", None), ("FROM ingestion.result_submission_grants", None), ("RETURNING expires_at", (expires_at,)), ] ) value = PostgresDirectIngestionRepository._issue_grant( cursor, JOB_ID, 1, 600, TOKEN, ) self.assertEqual(value.expires_at, expires_at) self.assertNotIn(TOKEN, repr(value)) grant_insert = next( item for item in cursor.executed if item[0].startswith( "INSERT INTO ingestion.result_submission_grants" ) ) assert grant_insert[1] is not None self.assertEqual( grant_insert[1][0], hashlib.sha256(TOKEN.encode()).hexdigest(), ) self.assertNotIn(TOKEN, grant_insert[1]) mode_update = next( query for query, _params in cursor.executed if query.startswith("UPDATE ingestion.processing_runs") ) self.assertIn("result_delivery_mode = 'direct_mcp'", mode_update) def test_terminal_exact_replay_returns_receipt_without_revalidation(self) -> None: value = request() now = datetime.now(timezone.utc) grant_row = ( 21, now + timedelta(minutes=5), now, None, 11, JOB_ID, "opera_daily", "accepted", "direct_mcp", 8, value.processor_version, value.rule_set_sha256, 12, 1, "succeeded", "opera_xml", SOURCE_KEY, "source.xml", SOURCE_HASH, 123, "application/xml", now, ) submission_row = ( 31, value.submission_key, "committed", DIRECT_CONTRACT_VERSION, value.business_date, value.processor_version, value.rule_set_sha256, value.result_schema_version, None, value.payload_sha256, len(value.payload_bytes), value.record_count, receipt(value), None, ) cursor = ScriptedCursor( [ ( "FROM ingestion.result_submission_grants AS submission_grant", grant_row, ), ("FROM ingestion.result_submissions", submission_row), ] ) state = PostgresDirectIngestionRepository._receive(cursor, value) self.assertEqual(state.status, "committed") self.assertIsNotNone(state.receipt) assert state.receipt is not None self.assertEqual(state.receipt.daily_version_id, 71) self.assertFalse( any( query.startswith("UPDATE ") or query.startswith("INSERT ") for query, _params in cursor.executed ) ) receive_query = cursor.executed[0][0] self.assertIn( "AS submission_grant", receive_query, ) self.assertNotIn(" AS grant", receive_query) def test_consumed_attempt_with_changed_payload_conflicts(self) -> None: original = request() changed = request(marker="different") now = datetime.now(timezone.utc) grant_row = ( 21, now + timedelta(minutes=5), now, None, 11, JOB_ID, "opera_daily", "validating", "direct_mcp", 8, original.processor_version, original.rule_set_sha256, 12, 1, "delivered", "opera_xml", SOURCE_KEY, "source.xml", SOURCE_HASH, 123, "application/xml", now, ) submission_row = ( 31, original.submission_key, "validating", DIRECT_CONTRACT_VERSION, original.business_date, original.processor_version, original.rule_set_sha256, original.result_schema_version, original.payload, original.payload_sha256, len(original.payload_bytes), original.record_count, {}, None, ) cursor = ScriptedCursor( [ ( "FROM ingestion.result_submission_grants AS submission_grant", grant_row, ), ("FROM ingestion.result_submissions", submission_row), ] ) with self.assertRaises(IngestionError) as raised: PostgresDirectIngestionRepository._receive(cursor, changed) self.assertEqual(raised.exception.code, "SUBMISSION_CONFLICT") def test_grant_ttl_and_token_factory_fail_closed(self) -> None: repository = PostgresDirectIngestionRepository( config=DatabaseConfig("postgresql://synthetic"), token_factory=lambda: "too-short", ) with self.assertRaises(IngestionError) as raised: repository.issue_grant(JOB_ID, 1, ttl_seconds=600) self.assertEqual(raised.exception.code, "SUBMISSION_GRANT_UNAVAILABLE") with self.assertRaises(IngestionError): repository.issue_grant(JOB_ID, 1, ttl_seconds=1801) if __name__ == "__main__": unittest.main(verbosity=2)