from __future__ import annotations import os import tempfile import unittest from pathlib import Path from unittest.mock import patch from arr_ingestion.contracts import IngestionError from arr_ingestion.postgres import ( DatabaseConfig, PostgresIngestionRepository, ) 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 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 PostgresIngestionTests(unittest.TestCase): 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_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) 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)