"""Unit tests for DB connection resilience + observability. These are DB-free: they exercise the application_name checkout hook against a fake DBAPI connection, the deriver polling backoff math, and the in-flight gauge listeners directly. """ from types import SimpleNamespace from typing import Any import pytest from sqlalchemy.ext.asyncio import AsyncSession import src.db as db_module from src.config import settings from src.db import DBQueryInflightTracker from src.telemetry.prometheus.metrics import db_queries_in_flight_gauge def test_session_local_uses_vanilla_async_session() -> None: """Regression guard: no custom session subclass / acquisition logic. Connection acquisition is a single lazy checkout owned by AsyncSession; there must be no re-introduced eager-checkout or retry hooks on the session. """ session = db_module.SessionLocal() assert type(session) is AsyncSession assert not hasattr(session, "_ensure_acquired") assert not hasattr(session, "_honcho_acquired") # --- application_name checkout hook ------------------------------------------ class _FakeCursor: def __init__(self, recorder: list[Any], raise_exc: Exception | None) -> None: self.recorder: list[Any] = recorder self.raise_exc: Exception | None = raise_exc self.closed: bool = False def execute(self, sql: str, params: Any = None) -> None: if self.raise_exc is not None: raise self.raise_exc self.recorder.append((sql, params)) def close(self) -> None: self.closed = True class _FakeDBAPIConn: def __init__(self, recorder: list[Any], raise_exc: Exception | None = None) -> None: self._cursor: _FakeCursor = _FakeCursor(recorder, raise_exc) def cursor(self) -> _FakeCursor: return self._cursor def test_checkout_hook_sets_application_name_from_request_context() -> None: recorder: list[Any] = [] conn = _FakeDBAPIConn(recorder) token = db_module.request_context.set("request:trace-ctx") try: db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] finally: db_module.request_context.reset(token) assert len(recorder) == 1 sql, params = recorder[0] assert "set_config" in sql and "application_name" in sql assert params == ("request:trace-ctx",) assert conn._cursor.closed is True # pyright: ignore[reportPrivateUsage] def test_checkout_hook_defaults_to_unknown_without_context() -> None: recorder: list[Any] = [] conn = _FakeDBAPIConn(recorder) token = db_module.request_context.set(None) try: db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] finally: db_module.request_context.reset(token) assert recorder[0][1] == ("unknown",) def test_checkout_hook_swallows_errors() -> None: """A failure tagging the connection must never break the checkout.""" conn = _FakeDBAPIConn([], raise_exc=RuntimeError("boom")) # Must not raise. db_module._set_application_name_on_checkout(conn, None, None) # pyright: ignore[reportPrivateUsage] # --- deriver polling backoff math -------------------------------------------- def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0) # Disable jitter so the schedule is asserted exactly (jitter is tested # separately in tests/deriver/test_queue_processing.py::TestPollingJitter). monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) from src.deriver.queue_manager import QueueManager qm = QueueManager() seq = [qm._advance_poll_interval() for _ in range(8)] # pyright: ignore[reportPrivateUsage] assert seq == [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0, 30.0] # caps at max qm._reset_poll_interval() # pyright: ignore[reportPrivateUsage] assert qm._advance_poll_interval() == 1.0 # pyright: ignore[reportPrivateUsage] def test_polling_backoff_disabled_stays_constant( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", False) monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) from src.deriver.queue_manager import QueueManager qm = QueueManager() assert [qm._advance_poll_interval() for _ in range(3)] == [1.0, 1.0, 1.0] # pyright: ignore[reportPrivateUsage] @pytest.mark.asyncio async def test_polling_loop_idle_sleeps_once_per_cycle( monkeypatch: pytest.MonkeyPatch, ) -> None: """Drive the real loop on an empty queue: exactly one (growing, capped) sleep per empty poll — no double-sleep from the queue_empty_flag branch.""" monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 8.0) monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0) import asyncio from src.deriver import queue_manager as qm_mod qm = qm_mod.QueueManager() sleeps: list[float] = [] polls = {"n": 0} async def fake_cleanup() -> None: return None async def fake_claim() -> dict[str, str]: polls["n"] += 1 if polls["n"] >= 5: qm.shutdown_event.set() # stop after 5 empty polls return {} async def fake_sleep(seconds: float) -> None: sleeps.append(seconds) monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) # queue_manager calls asyncio.sleep on the stdlib module; patch it there. monkeypatch.setattr(asyncio, "sleep", fake_sleep) await qm.polling_loop() # One sleep per empty poll, growing 1->2->4->8 then capped at 8 (not doubled). assert sleeps == [1.0, 2.0, 4.0, 8.0, 8.0] def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test") child: Any = db_queries_in_flight_gauge.labels(instance_type="api") tracker = DBQueryInflightTracker(child) key = DBQueryInflightTracker.INFLIGHT_KEY def value() -> float: return float(child._value.get()) start = value() conn = SimpleNamespace(info={}) # Normal execute: before -> after returns to baseline. tracker.on_before(conn) assert value() == start + 1 assert conn.info[key] is True tracker.on_after(conn) assert value() == start assert key not in conn.info # Errored execute: before -> on_error decrements (after never fires). tracker.on_before(conn) assert value() == start + 1 tracker.on_error(SimpleNamespace(connection=conn)) assert value() == start # on_error without a matching before (e.g. connect error) must not push the # gauge negative. tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={}))) assert value() == start