next-tjwater-server/tests/unit/test_realtime_repository.py

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)