yuxuanhui
46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
"""Lifecycle tests for the shared selection PostgreSQL pool owner."""
|
|
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
|
|
from zhixing_server.modules.selection.infrastructure.postgres_pool import SelectionPostgresPool
|
|
|
|
|
|
class FakePool:
|
|
def __init__(self) -> None:
|
|
self.open_calls = 0
|
|
self.close_calls = 0
|
|
self.connection_calls = 0
|
|
|
|
def open(self, *, wait: bool = True) -> None:
|
|
assert wait is True
|
|
self.open_calls += 1
|
|
|
|
def close(self) -> None:
|
|
self.close_calls += 1
|
|
|
|
@contextmanager
|
|
def connection(self) -> Generator[str, None, None]:
|
|
self.connection_calls += 1
|
|
yield "connection"
|
|
|
|
|
|
def test_selection_pool_opens_once_borrows_and_closes_injected_pool() -> None:
|
|
fake = FakePool()
|
|
owner = SelectionPostgresPool(
|
|
"postgresql://test",
|
|
max_connections=6,
|
|
pool=fake,
|
|
)
|
|
|
|
with owner.connection() as connection:
|
|
assert connection == "connection"
|
|
with owner.connection() as connection:
|
|
assert connection == "connection"
|
|
|
|
assert owner.max_connections == 6
|
|
assert fake.open_calls == 1
|
|
assert fake.connection_calls == 2
|
|
owner.close()
|
|
assert fake.close_calls == 1
|