300 lines
9.2 KiB
Python
300 lines
9.2 KiB
Python
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)
|