Files
wyndham-ARR/tests/test_arr_ingestion_postgres.py
2026-07-31 15:11:42 +08:00

414 lines
15 KiB
Python

from __future__ import annotations
import os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from arr_ingestion.contracts import ArtifactRef, IngestionError
from arr_ingestion.postgres import (
DatabaseConfig,
PostgresIngestionRepository,
)
from arr_ingestion.repository import JobRegistration
from tests.test_arr_opera_daily_ingest import run_processor, success_xml
class InsertCursor:
def __init__(self) -> None:
self.params: list[tuple[object, ...]] = []
self._next_id = 100
def execute(self, query: str, params: tuple[object, ...]) -> None:
self.assert_placeholder_count(query, params)
self.params.append(params)
@staticmethod
def assert_placeholder_count(query: str, params: tuple[object, ...]) -> None:
if query.count("%s") != len(params):
raise AssertionError(
f"placeholder mismatch: {query.count('%s')} != {len(params)}"
)
def fetchone(self) -> tuple[int]:
current = self._next_id
self._next_id += 1
return (current,)
class BookingCursor:
def __init__(self) -> None:
self.executed: tuple[str, tuple[object, ...]] | None = None
def execute(self, query: str, params: tuple[object, ...]) -> None:
InsertCursor.assert_placeholder_count(query, params)
self.executed = (query, params)
def fetchall(self):
return [
("GROUP-MATCH", 1),
("GROUP-REPEATED", 2),
]
class ArtifactCursor:
def __init__(self, existing=None, inserted_id: int = 101) -> None:
self.existing = existing
self.inserted_id = inserted_id
self.last_query = ""
self.executed: list[tuple[str, tuple[object, ...]]] = []
def execute(self, query: str, params: tuple[object, ...]) -> None:
InsertCursor.assert_placeholder_count(query, params)
self.last_query = " ".join(query.split())
self.executed.append((self.last_query, params))
def fetchone(self):
if self.last_query.startswith("SELECT"):
return self.existing
if self.last_query.startswith("INSERT INTO ingestion.artifacts"):
return (self.inserted_id,)
return None
class SqlStateError(RuntimeError):
def __init__(self, sqlstate: str):
super().__init__("synthetic database error")
self.sqlstate = sqlstate
class RetryConnection:
def __init__(self) -> None:
self.closed = False
def transaction(self):
return self
def cursor(self):
return self
def __enter__(self):
return self
def __exit__(self, _kind, _value, _traceback):
return False
def close(self) -> None:
self.closed = True
class LifecycleCursor:
def __init__(self, run_status: str = "queued", attempt_status: str = "queued") -> None:
self.run_status = run_status
self.attempt_status = attempt_status
self.last_query = ""
self.executed: list[tuple[str, tuple[object, ...]]] = []
def execute(self, query: str, params: tuple[object, ...]) -> None:
InsertCursor.assert_placeholder_count(query, params)
self.last_query = " ".join(query.split())
self.executed.append((self.last_query, params))
def fetchone(self):
if self.last_query.startswith("SELECT run.id, run.run_status"):
return (12, self.run_status, 13, self.attempt_status)
return None
class RegistrationCursor:
def __init__(self) -> None:
self.last_query = ""
self.executed: list[tuple[str, tuple[object, ...]]] = []
def execute(self, query: str, params: tuple[object, ...]) -> None:
InsertCursor.assert_placeholder_count(query, params)
self.last_query = " ".join(query.split())
self.executed.append((self.last_query, params))
def fetchone(self):
if self.last_query.startswith(
"SELECT id, artifact_kind, original_filename"
):
return None
if self.last_query.startswith("INSERT INTO ingestion.artifacts"):
return (101,)
if self.last_query.startswith(
"SELECT id, source_artifact_id, requested_processor_version"
):
return None
if self.last_query.startswith("INSERT INTO ingestion.processing_runs"):
return (202,)
return None
class PostgresIngestionTests(unittest.TestCase):
@staticmethod
def artifact(object_key: str, sha256: str = "a" * 64) -> ArtifactRef:
return ArtifactRef.from_dict(
"source_xml",
{
"object_key": object_key,
"original_filename": "source.xml",
"sha256": sha256,
"byte_size": 321,
"mime_type": "application/xml",
},
)
def test_database_config_is_environment_only(self):
with patch.dict(os.environ, {}, clear=True):
with self.assertRaisesRegex(IngestionError, "ARR_DATABASE_URL"):
DatabaseConfig.from_environment()
with patch.dict(os.environ, {"ARR_DATABASE_URL": "postgresql://synthetic"}, clear=True):
self.assertEqual(
DatabaseConfig.from_environment().dsn,
"postgresql://synthetic",
)
def test_record_insert_maps_duplicate_lineage_and_booking_source_status(self):
with tempfile.TemporaryDirectory() as temp_dir:
exit_code, _xml, _output, _result, payload = run_processor(
success_xml(), Path(temp_dir)
)
self.assertEqual(exit_code, 0)
cursor = InsertCursor()
links = {
"AB-123": ("matched", 1),
"GO-EASY-GROUP": ("matched", 2),
}
PostgresIngestionRepository._insert_records(
cursor,
9,
payload["records"],
links,
)
self.assertEqual(len(cursor.params), payload["source_rows"])
for params in cursor.params:
self.assertEqual(params[0], 9)
self.assertEqual(len(params), 33)
duplicate_params = cursor.params[2]
self.assertEqual(duplicate_params[7], 101)
matched_params = cursor.params[1]
self.assertEqual(matched_params[-2:], ("matched", 1))
missing_group_params = cursor.params[3]
self.assertEqual(
missing_group_params[-2:], ("missing_group_code", 0)
)
repeated_params = cursor.params[4]
self.assertEqual(repeated_params[-2:], ("matched", 2))
def test_booking_lookup_accepts_repeated_group_rows_and_marks_unmatched(self):
records = [
{"group_code_key": "GROUP-MATCH"},
{"group_code_key": "GROUP-REPEATED"},
{"group_code_key": "GROUP-MISSING"},
{"group_code_key": None},
]
cursor = BookingCursor()
links = PostgresIngestionRepository._load_booking_links(cursor, records)
self.assertEqual(links["GROUP-MATCH"], ("matched", 1))
self.assertEqual(links["GROUP-REPEATED"], ("matched", 2))
self.assertEqual(links["GROUP-MISSING"], ("unmatched", 0))
self.assertIsNotNone(cursor.executed)
def test_artifact_identity_is_storage_object_not_content_hash(self):
key = (
"arr/jobs/arrjob-repeat/attempts/0001/committed/"
"source_xml/source.xml"
)
cursor = ArtifactCursor()
artifact_id = PostgresIngestionRepository._ensure_artifact(
cursor,
self.artifact(key),
)
self.assertEqual(artifact_id, 101)
select_query, select_params = cursor.executed[0]
self.assertIn("object_key = %s", select_query)
self.assertIn("object_version_id IS NULL", select_query)
self.assertNotIn("WHERE artifact_kind = %s", select_query)
self.assertEqual(select_params, ("oss", "arr-private", key))
self.assertTrue(
any(query.startswith("INSERT INTO ingestion.artifacts") for query, _ in cursor.executed)
)
def test_existing_storage_identity_is_reused_only_when_all_metadata_matches(self):
key = (
"arr/jobs/arrjob-existing/attempts/0001/committed/"
"source_xml/source.xml"
)
cursor = ArtifactCursor(
(77, "opera_xml", "source.xml", "a" * 64, 321, "application/xml")
)
artifact_id = PostgresIngestionRepository._ensure_artifact(
cursor,
self.artifact(key),
)
self.assertEqual(artifact_id, 77)
self.assertEqual(len(cursor.executed), 1)
def test_existing_storage_identity_rejects_conflicting_content(self):
key = (
"arr/jobs/arrjob-conflict/attempts/0001/committed/"
"source_xml/source.xml"
)
cursor = ArtifactCursor(
(77, "opera_xml", "source.xml", "b" * 64, 321, "application/xml")
)
with self.assertRaises(IngestionError) as raised:
PostgresIngestionRepository._ensure_artifact(
cursor,
self.artifact(key),
)
self.assertEqual(raised.exception.code, "ARTIFACT_CONFLICT")
def test_job_registration_persists_uploaded_filename_separately(self):
cursor = RegistrationCursor()
registration = JobRegistration(
job_id="arrjob-upload-name-001",
source=self.artifact(
"arr/jobs/arrjob-upload-name-001/attempts/0001/"
"committed/source_xml/source.xml"
),
processor_version="3.0.0",
rule_set_sha256="b" * 64,
attempt_no=1,
idempotency_key="c" * 64,
uploaded_filename="0723-Opera.XML",
)
repository = PostgresIngestionRepository(
DatabaseConfig("postgresql://synthetic")
)
repository._register_job(cursor, registration)
run_insert = next(
(query, params)
for query, params in cursor.executed
if query.startswith("INSERT INTO ingestion.processing_runs")
)
self.assertIn("uploaded_filename", run_insert[0])
self.assertEqual(run_insert[1][-1], "0723-Opera.XML")
artifact_insert = next(
(query, params)
for query, params in cursor.executed
if query.startswith("INSERT INTO ingestion.artifacts")
)
self.assertEqual(artifact_insert[1][4], "source.xml")
def test_transaction_reopens_after_serialization_failure(self):
connections = [RetryConnection(), RetryConnection()]
connect_count = 0
operation_count = 0
def connect(_dsn: str):
nonlocal connect_count
connection = connections[connect_count]
connect_count += 1
return connection
def operation(_cursor):
nonlocal operation_count
operation_count += 1
if operation_count == 1:
raise SqlStateError("40001")
return "committed"
repository = PostgresIngestionRepository(
DatabaseConfig("postgresql://synthetic"), connect=connect
)
with patch.object(repository, "_begin"), patch("arr_ingestion.postgres.time.sleep"):
result = repository._run_transaction(operation, "safe failure")
self.assertEqual(result, "committed")
self.assertEqual(connect_count, 2)
self.assertEqual(operation_count, 2)
self.assertTrue(all(connection.closed for connection in connections))
def test_nontransient_database_failure_is_not_retried(self):
connection = RetryConnection()
repository = PostgresIngestionRepository(
DatabaseConfig("postgresql://synthetic"), connect=lambda _dsn: connection
)
with patch.object(repository, "_begin"):
with self.assertRaisesRegex(IngestionError, "safe failure") as raised:
repository._run_transaction(
lambda _cursor: (_ for _ in ()).throw(SqlStateError("22000")),
"safe failure",
)
self.assertEqual(raised.exception.code, "DATABASE_WRITE_FAILED")
self.assertTrue(connection.closed)
def test_programmatic_start_marks_attempt_and_run_running(self):
cursor = LifecycleCursor()
repository = PostgresIngestionRepository(DatabaseConfig("postgresql://synthetic"))
repository._mark_running(
cursor,
"arrjob-programmatic-001",
1,
)
statements = "\n".join(query for query, _params in cursor.executed)
self.assertIn("attempt_status = 'running'", statements)
self.assertIn("run_status = 'running'", statements)
self.assertIn("started_at = COALESCE(started_at, now())", statements)
def test_programmatic_infrastructure_failure_is_terminal_and_outboxed(self):
cursor = LifecycleCursor(run_status="running", attempt_status="running")
repository = PostgresIngestionRepository(DatabaseConfig("postgresql://synthetic"))
repository._record_runtime_failure(
cursor,
"arrjob-programmatic-002",
1,
"PROCESSOR_TIMEOUT",
)
statements = "\n".join(query for query, _params in cursor.executed)
self.assertIn("attempt_status = 'failed'", statements)
self.assertIn("run_status = 'failed'", statements)
self.assertIn("INSERT INTO ingestion.outbox_events", statements)
rendered_params = repr([params for _query, params in cursor.executed])
self.assertIn("PROCESSOR_TIMEOUT", rendered_params)
self.assertNotIn("password", rendered_params.lower())
class Migration008ContractTests(unittest.TestCase):
def test_rebuild_targets_v3_and_all_source_record_storage(self):
project_root = Path(__file__).resolve().parents[1]
sql = (project_root / "database" / "008_arr_mvp_v1_rebuild.sql").read_text(
encoding="utf-8"
)
for required in (
"CREATE TABLE ingestion.processing_runs",
"CREATE TABLE ingestion.processing_attempts",
"CREATE TABLE ingestion.processing_deliveries",
"CREATE TABLE ingestion.outbox_events",
"structured_result_artifact_id",
"CREATE TABLE booking.source_rows",
"CREATE TABLE booking.parse_versions",
"CREATE TABLE finance.daily_records",
"CREATE VIEW finance.v_monthly_report_rows",
"CREATE VIEW finance.v_channel_details",
"outcome <> 'retained'",
"total_price = real_price * no_of_rooms * nights",
):
with self.subTest(required=required):
self.assertIn(required, sql)
self.assertNotIn("CREATE TABLE finance.report_versions", sql)
self.assertNotIn("CREATE TABLE booking.file_objects", sql)
if __name__ == "__main__":
unittest.main(verbosity=2)