import asyncio from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone import pytest from pydantic import ValidationError from app.core.config import Settings, settings from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository SNAPSHOT_BUDGET_SQL = "SELECT set_config('timescaledb.max_tuples_decompressed_per_dml_transaction', %s, true)" class _FakeCursor: def __init__(self): self.calls: list[tuple[str, tuple]] = [] async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): return False async def execute(self, query, params): self.calls.append((str(query), params)) async def fetchall(self): return [] class _FakeConnection: def __init__(self): self.cursor_instance = _FakeCursor() def cursor(self): return self.cursor_instance def test_get_links_by_time_range_normalizes_inputs_to_utc(): conn = _FakeConnection() asyncio.run( RealtimeRepository.get_links_by_time_range( conn, datetime.fromisoformat("2026-06-01T08:00:00+08:00"), datetime.fromisoformat("2026-06-01T09:00:00+08:00"), ) ) assert len(conn.cursor_instance.calls) == 1 _, params = conn.cursor_instance.calls[0] assert params == ( datetime(2026, 6, 1, 0, 0, tzinfo=timezone.utc), datetime(2026, 6, 1, 1, 0, tzinfo=timezone.utc), ) def test_get_nodes_by_time_range_normalizes_inputs_to_utc(): conn = _FakeConnection() asyncio.run( RealtimeRepository.get_nodes_by_time_range( conn, datetime.fromisoformat("2026-06-01T08:00:00+08:00"), datetime.fromisoformat("2026-06-01T09:00:00+08:00"), ) ) assert len(conn.cursor_instance.calls) == 1 _, params = conn.cursor_instance.calls[0] assert params == ( datetime(2026, 6, 1, 0, 0, tzinfo=timezone.utc), datetime(2026, 6, 1, 1, 0, tzinfo=timezone.utc), ) class _SyncTransactionConnection: def __init__(self): self.transactions = 0 self.calls: list[tuple[str, tuple]] = [] @contextmanager def transaction(self): self.transactions += 1 yield @contextmanager def cursor(self): connection = self class Cursor: def execute(self, query, params): connection.calls.append((query, params)) yield Cursor() def test_realtime_node_and_link_replacement_share_outer_transaction(monkeypatch): conn = _SyncTransactionConnection() calls: list[str] = [] monkeypatch.setattr( RealtimeRepository, "_copy_nodes_sync", lambda _cur, _data, _time: calls.append("nodes"), ) monkeypatch.setattr( RealtimeRepository, "_copy_links_sync", lambda _cur, _data, _time: calls.append("links"), ) RealtimeRepository.store_realtime_simulation_result_sync( conn, [{"node": "N1", "result": [{"pressure": 1.0}]}], [{"link": "L1", "result": [{"flow": 2.0}]}], "2026-06-01T00:00:00Z", ) assert conn.transactions == 1 assert calls == ["nodes", "links"] assert [query for query, _params in conn.calls] == [ SNAPSHOT_BUDGET_SQL, "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))", "DELETE FROM realtime.node_results WHERE time = %s", "DELETE FROM realtime.link_results WHERE time = %s", ] assert conn.calls[0][1] == (str(settings.TIMESCALEDB_REALTIME_MAX_DECOMPRESSED_TUPLES),) def test_realtime_batch_rejects_multiple_timestamps_before_writing(): data = [ {"time": "2026-06-01T00:00:00Z", "id": "N1"}, {"time": "2026-06-01T00:15:00Z", "id": "N2"}, ] try: RealtimeRepository.insert_nodes_batch_sync(object(), data) except ValueError as exc: assert str(exc) == "Realtime batch must contain exactly one timestamp" else: raise AssertionError("multiple realtime timestamps were accepted") class _AsyncEmptySnapshotConnection: def __init__(self): self.calls: list[tuple[str, tuple]] = [] @asynccontextmanager async def transaction(self): yield def cursor(self): connection = self class Cursor: async def __aenter__(self): return self async def __aexit__(self, *_args): return False async def execute(self, query, params): connection.calls.append((query, params)) return Cursor() def test_empty_realtime_side_is_deleted_as_part_of_snapshot_replacement(): conn = _AsyncEmptySnapshotConnection() asyncio.run( RealtimeRepository.store_realtime_simulation_result( conn, node_result_list=[], link_result_list=[], result_start_time="2026-06-01T00:00:00Z", ) ) assert [query for query, _params in conn.calls] == [ SNAPSHOT_BUDGET_SQL, "SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))", "DELETE FROM realtime.node_results WHERE time = %s", "DELETE FROM realtime.link_results WHERE time = %s", ] @pytest.mark.parametrize("side", ["nodes", "links"]) @pytest.mark.parametrize("asynchronous", [False, True]) def test_all_snapshot_batch_writers_use_configured_local_budget(monkeypatch, side, asynchronous): monkeypatch.setattr(settings, "TIMESCALEDB_REALTIME_MAX_DECOMPRESSED_TUPLES", 1_200_000) conn = _AsyncEmptySnapshotConnection() if asynchronous else _SyncTransactionConnection() if asynchronous: async def copy(*_): pass monkeypatch.setattr(RealtimeRepository, f"_copy_{side}", copy) else: monkeypatch.setattr(RealtimeRepository, f"_copy_{side}_sync", lambda *_: None) method = getattr(RealtimeRepository, f"insert_{side}_batch" + ("" if asynchronous else "_sync")) rows = [{"time": "2026-10-08T04:00:00Z", "id": "fixture"}] if asynchronous: asyncio.run(method(conn, rows)) else: method(conn, rows) assert conn.calls[0] == (SNAPSHOT_BUDGET_SQL, ("1200000",)) assert "pg_advisory_xact_lock" in conn.calls[1][0] assert conn.calls[2][0].startswith("DELETE FROM realtime.") @pytest.mark.parametrize("value", [0, -1, 2_147_483_648]) def test_snapshot_budget_must_remain_a_finite_positive_postgres_limit(value): with pytest.raises(ValidationError): Settings(TIMESCALEDB_REALTIME_MAX_DECOMPRESSED_TUPLES=value, _env_file=None)