dzentra_bot/app/tests/unit/storage/test_postgres_pool.py

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]