"""PostgreSQL auto-reconcile transaction sequencing tests.""" from __future__ import annotations from datetime import UTC, datetime from types import SimpleNamespace from uuid import uuid4 import psycopg import pytest from trex_core.adapters.persistence.auto_reconcile import PostgreSQLAutoReconcileStore from trex_core.adapters.persistence.postgresql import PostgreSQLSlotRepository from trex_core.adapters.portfolio.reconcile_warning_store import ( PostgreSQLReconcileWarningStore, ) from trex_core.application.auto_reconcile import AutoReconcileCommand from trex_core.domain.slot import SlotService, SlotState class FakeCursor: def __init__(self, rowcounts=(1, 1), fail_on_insert=False): self.rowcounts = iter(rowcounts) self.rowcount = -1 self.fail_on_insert = fail_on_insert self.calls = [] def __enter__(self): return self def __exit__(self, exc_type, exc, tb): return False def execute(self, sql, params=None): normalized = " ".join(sql.split()) self.calls.append((normalized, params)) if normalized.startswith("UPDATE daily_slot"): self.rowcount = next(self.rowcounts) elif normalized.startswith("DELETE FROM slot_reconcile_warning"): self.rowcount = next(self.rowcounts) elif normalized.startswith("INSERT INTO slot_transition_log") and self.fail_on_insert: raise RuntimeError("audit insert failed") class FakeConnection: def __init__(self, cursor): self._cursor = cursor self.commits = 0 self.rollbacks = 0 def cursor(self): return self._cursor def commit(self): self.commits += 1 def rollback(self): self.rollbacks += 1 def close(self): pass class FakeConnectionContext: def __init__(self, conn): self.conn = conn def __enter__(self): return self.conn def __exit__(self, exc_type, exc, tb): return False def _command(): return AutoReconcileCommand( slot_id=uuid4(), instance_key="inst-a", portfolio_id="pf-a", symbol="019685", expected_qty=1452.0, expected_avg_price=2100.0, target_qty=726.0, target_avg_price=2172.46829187, hub_quantity=726.0, broker_quantity=726, evidence_observed_at=datetime.now(tz=UTC), ) def _install(monkeypatch, conn): module = SimpleNamespace(connect=lambda dsn, autocommit=False: FakeConnectionContext(conn)) monkeypatch.setattr( "trex_core.adapters.persistence.auto_reconcile._require_psycopg", lambda _module=None: module, ) def test_postgresql_auto_reconcile_commits_update_b_clear_and_audit(monkeypatch): cursor = FakeCursor(rowcounts=(1, 1)) conn = FakeConnection(cursor) _install(monkeypatch, conn) assert PostgreSQLAutoReconcileStore("postgresql://unused").apply_holding_mirror(_command()) assert conn.commits == 1 assert conn.rollbacks == 0 statements = [sql for sql, _params in cursor.calls] assert statements[0].startswith("SELECT pg_advisory_xact_lock") assert "status='holding'" in statements[1] assert "avg_price IS NOT DISTINCT FROM" in statements[1] assert "broker_order_id IS NULL" in statements[1] assert statements[2].startswith("DELETE FROM slot_reconcile_warning") assert cursor.calls[2][1][1] == "B" assert statements[3].startswith("INSERT INTO slot_transition_log") assert cursor.calls[3][1][2:4] == ( "hub_broker_matched_core_mirror_stale", "auto_reconcile", ) def test_postgresql_auto_reconcile_cas_failure_rolls_back_without_warning_or_audit(monkeypatch): cursor = FakeCursor(rowcounts=(0,)) conn = FakeConnection(cursor) _install(monkeypatch, conn) assert not PostgreSQLAutoReconcileStore("postgresql://unused").apply_holding_mirror(_command()) assert conn.commits == 0 assert conn.rollbacks == 1 assert len(cursor.calls) == 2 def test_postgresql_auto_reconcile_audit_failure_rolls_back_slot_and_warning(monkeypatch): cursor = FakeCursor(rowcounts=(1, 1), fail_on_insert=True) conn = FakeConnection(cursor) _install(monkeypatch, conn) with pytest.raises(RuntimeError, match="audit insert failed"): PostgreSQLAutoReconcileStore("postgresql://unused").apply_holding_mirror(_command()) assert conn.commits == 0 assert conn.rollbacks == 1 assert len(cursor.calls) == 4 def _seed_real_transaction_state(pg_clean): repo = PostgreSQLSlotRepository(pg_clean) service = SlotService(repo, instance_key="inst-a") warning = PostgreSQLReconcileWarningStore(pg_clean) repo.reconcile_universe("inst-a", {"pf-a": "019685"}) slot = repo.find_slot_by_portfolio("inst-a", "pf-a") assert slot is not None pending = service.transit_slot( slot.slot_id, SlotState.INACTIVE, SlotState.PENDING_BUY, reason="buy", actor="test", symbol="019685", ) holding = service.transit_slot( pending.slot_id, SlotState.PENDING_BUY, SlotState.HOLDING, reason="filled", actor="test", qty=1452.0, avg_price=2100.0, ) warning.set_warned( str(holding.slot_id), portfolio_id="pf-a", case_code="B", actor="reconcile", ) warning.set_warned( str(holding.slot_id), portfolio_id="pf-a", case_code="C", actor="reconcile", ) command = AutoReconcileCommand( slot_id=holding.slot_id, instance_key="inst-a", portfolio_id="pf-a", symbol="019685", expected_qty=1452.0, expected_avg_price=2100.0, target_qty=726.0, target_avg_price=2172.46829187, hub_quantity=726.0, broker_quantity=726, evidence_observed_at=datetime.now(tz=UTC), ) return repo, warning, holding, command def test_postgresql_auto_reconcile_real_transaction_commit(pg_clean): repo, warning, holding, command = _seed_real_transaction_state(pg_clean) assert PostgreSQLAutoReconcileStore(pg_clean).apply_holding_mirror(command) current = repo.fetch_slot(holding.slot_id) assert current is not None assert current.status is SlotState.HOLDING assert current.qty == 726.0 assert current.avg_price == pytest.approx(2172.46829187) info = warning.get_warning_info(str(holding.slot_id)) assert info is not None assert info["case_codes"] == ["C"] with psycopg.connect(pg_clean) as conn: row = conn.execute( "SELECT actor, reason FROM slot_transition_log " "WHERE slot_id = %s ORDER BY ts DESC LIMIT 1", (holding.slot_id,), ).fetchone() assert row == ("auto_reconcile", "hub_broker_matched_core_mirror_stale") def test_postgresql_auto_reconcile_real_transaction_audit_failure_rolls_back(pg_clean): repo, warning, holding, command = _seed_real_transaction_state(pg_clean) with psycopg.connect(pg_clean) as conn: before_audit_row = conn.execute( "SELECT count(*) FROM slot_transition_log WHERE slot_id = %s", (holding.slot_id,), ).fetchone() assert before_audit_row is not None before_audit_count = before_audit_row[0] with psycopg.connect(pg_clean, autocommit=True) as conn: conn.execute( """ CREATE FUNCTION trex_test_fail_auto_reconcile_audit() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'forced auto-reconcile audit failure'; END $$ """ ) conn.execute( """ CREATE TRIGGER trex_test_fail_auto_reconcile_audit BEFORE INSERT ON slot_transition_log FOR EACH ROW WHEN (NEW.actor = 'auto_reconcile') EXECUTE FUNCTION trex_test_fail_auto_reconcile_audit() """ ) try: with pytest.raises(psycopg.errors.RaiseException, match="forced auto-reconcile"): PostgreSQLAutoReconcileStore(pg_clean).apply_holding_mirror(command) current = repo.fetch_slot(holding.slot_id) assert current is not None assert current.qty == 1452.0 assert current.avg_price == 2100.0 info = warning.get_warning_info(str(holding.slot_id)) assert info is not None assert info["case_codes"] == ["B", "C"] with psycopg.connect(pg_clean) as conn: after_audit_row = conn.execute( "SELECT count(*) FROM slot_transition_log WHERE slot_id = %s", (holding.slot_id,), ).fetchone() assert after_audit_row is not None after_audit_count = after_audit_row[0] assert after_audit_count == before_audit_count finally: with psycopg.connect(pg_clean, autocommit=True) as conn: conn.execute( "DROP TRIGGER IF EXISTS trex_test_fail_auto_reconcile_audit ON slot_transition_log" ) conn.execute("DROP FUNCTION IF EXISTS trex_test_fail_auto_reconcile_audit()")