from __future__ import annotations from contextlib import nullcontext from typing import Any import pytest from src.storage.exceptions import PostgresConnectionPoolError from src.storage.postgres_pool import PostgresConnectionPool class RecordingPool: def __init__( self, *, open_error: BaseException | None = None, close_error: Exception | None = None, ) -> None: self.close_calls: list[float] = [] self.close_error = close_error self.connection_calls: list[float] = [] self.open_calls: list[tuple[bool, float]] = [] self.open_error = open_error self.connection_value = object() def open(self, *, wait: bool, timeout: float) -> None: self.open_calls.append((wait, timeout)) if self.open_error is not None: raise self.open_error def connection(self, *, timeout: float): self.connection_calls.append(timeout) return nullcontext(self.connection_value) def close(self, *, timeout: float) -> None: self.close_calls.append(timeout) if self.close_error is not None: raise self.close_error class RecordingPoolFactory: def __init__(self, pool: RecordingPool) -> None: self.calls: list[dict[str, Any]] = [] self.pool = pool def __call__(self, **kwargs: Any) -> RecordingPool: self.calls.append(kwargs) return self.pool def _connection_pool( *, pool: RecordingPool | None = None, ) -> tuple[ PostgresConnectionPool, RecordingPool, RecordingPoolFactory, ]: recording_pool = pool or RecordingPool() factory = RecordingPoolFactory(recording_pool) connection_pool = PostgresConnectionPool( conninfo="postgresql://db/dzentra", min_size=2, max_size=6, timeout_seconds=7.5, name="market-data", pool_factory=factory, ) return connection_pool, recording_pool, factory def test_constructor_does_not_create_or_open_pool() -> None: connection_pool, pool, factory = _connection_pool() assert connection_pool.is_open is False assert factory.calls == [] assert pool.open_calls == [] def test_open_creates_pool_with_explicit_configuration_and_waits() -> None: connection_pool, pool, factory = _connection_pool() connection_pool.open() assert connection_pool.is_open is True assert factory.calls == [ { "conninfo": "postgresql://db/dzentra", "min_size": 2, "max_size": 6, "timeout": 7.5, "kwargs": {"autocommit": False}, "name": "market-data", "open": False, } ] assert pool.open_calls == [(True, 7.5)] def test_repeated_open_is_no_op() -> None: connection_pool, pool, factory = _connection_pool() connection_pool.open() connection_pool.open() assert len(factory.calls) == 1 assert pool.open_calls == [(True, 7.5)] def test_connection_requires_open_pool() -> None: connection_pool, _, _ = _connection_pool() with pytest.raises(PostgresConnectionPoolError, match="not open"): connection_pool.connection() def test_connection_borrows_from_pool_with_timeout() -> None: connection_pool, pool, _ = _connection_pool() connection_pool.open() with connection_pool.connection() as connection: assert connection is pool.connection_value assert pool.connection_calls == [7.5] def test_close_is_idempotent() -> None: connection_pool, pool, _ = _connection_pool() connection_pool.open() connection_pool.close() connection_pool.close() assert connection_pool.is_open is False assert pool.close_calls == [7.5] def test_open_failure_closes_partial_pool_and_is_retryable() -> None: pool = RecordingPool(open_error=RuntimeError("database unavailable")) connection_pool, _, factory = _connection_pool(pool=pool) with pytest.raises(PostgresConnectionPoolError) as error_info: connection_pool.open() assert isinstance(error_info.value.__cause__, RuntimeError) assert connection_pool.is_open is False assert pool.close_calls == [7.5] pool.open_error = None connection_pool.open() assert len(factory.calls) == 2 assert connection_pool.is_open is True def test_open_failure_preserves_cleanup_failure_as_note() -> None: pool = RecordingPool( open_error=RuntimeError("open failed"), close_error=RuntimeError("close failed"), ) connection_pool, _, _ = _connection_pool(pool=pool) with pytest.raises(PostgresConnectionPoolError) as error_info: connection_pool.open() cause = error_info.value.__cause__ assert isinstance(cause, RuntimeError) assert cause.__notes__ == [ "PostgreSQL pool cleanup also failed: RuntimeError." ] def test_close_failure_leaves_wrapper_closed() -> None: pool = RecordingPool(close_error=RuntimeError("close failed")) connection_pool, _, _ = _connection_pool(pool=pool) connection_pool.open() with pytest.raises(PostgresConnectionPoolError) as error_info: connection_pool.close() assert isinstance(error_info.value.__cause__, RuntimeError) assert connection_pool.is_open is False connection_pool.close() assert pool.close_calls == [7.5] def test_keyboard_interrupt_is_not_wrapped() -> None: def interrupted_factory(**kwargs: object) -> object: raise KeyboardInterrupt connection_pool = PostgresConnectionPool( conninfo="postgresql://db/dzentra", pool_factory=interrupted_factory, ) with pytest.raises(KeyboardInterrupt): connection_pool.open() def test_keyboard_interrupt_during_pool_open_closes_partial_pool() -> None: pool = RecordingPool(open_error=KeyboardInterrupt()) connection_pool, _, _ = _connection_pool(pool=pool) with pytest.raises(KeyboardInterrupt): connection_pool.open() assert connection_pool.is_open is False assert pool.open_calls == [(True, 7.5)] assert pool.close_calls == [7.5] @pytest.mark.parametrize( ("override", "message"), ( ({"conninfo": ""}, "conninfo"), ({"min_size": 0}, "min_size"), ({"min_size": True}, "min_size"), ({"min_size": 2, "max_size": 1}, "max_size"), ({"timeout_seconds": 0}, "timeout_seconds"), ({"timeout_seconds": float("inf")}, "timeout_seconds"), ({"name": " "}, "name"), ), ) def test_constructor_rejects_invalid_configuration( override: dict[str, object], message: str, ) -> None: arguments: dict[str, object] = { "conninfo": "postgresql://db/dzentra", } arguments.update(override) with pytest.raises(ValueError, match=message): PostgresConnectionPool(**arguments) # type: ignore[arg-type]