import asyncio from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository 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] == [ "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", ] 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] == [ "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", ]