105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
import asyncio
|
|
from collections import OrderedDict
|
|
from contextlib import asynccontextmanager
|
|
from uuid import uuid4
|
|
|
|
from app.infra.db import dynamic_manager
|
|
from app.infra.db.timescaledb import sync_pool
|
|
from app.native.wndb.core import connection
|
|
|
|
|
|
class RecordingAsyncPool:
|
|
created: list[dict] = []
|
|
|
|
@staticmethod
|
|
async def check_connection(_conn) -> None:
|
|
return None
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
self.created.append(kwargs)
|
|
self.closed = False
|
|
|
|
async def open(self) -> None:
|
|
return None
|
|
|
|
async def close(self) -> None:
|
|
self.closed = True
|
|
|
|
@asynccontextmanager
|
|
async def connection(self):
|
|
yield object()
|
|
|
|
|
|
class RecordingPool:
|
|
created: list[dict] = []
|
|
|
|
@staticmethod
|
|
def check_connection(_conn) -> None:
|
|
return None
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
self.created.append(kwargs)
|
|
self.closed = False
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
def test_dynamic_project_pools_check_connections_before_borrow(monkeypatch) -> None:
|
|
async def create_pools() -> None:
|
|
manager = dynamic_manager.ProjectConnectionManager()
|
|
async with manager.pg_connection(
|
|
uuid4(), "biz_data", "postgresql://user:password@db.example/biz", 1, 5
|
|
):
|
|
pass
|
|
async with manager.timescale_connection(
|
|
uuid4(), "iot_data", "postgresql://user:password@db.example/ts", 1, 5
|
|
):
|
|
pass
|
|
|
|
RecordingAsyncPool.created = []
|
|
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", RecordingAsyncPool)
|
|
asyncio.run(create_pools())
|
|
|
|
assert len(RecordingAsyncPool.created) == 2
|
|
assert all(
|
|
options["check"] is dynamic_manager._check_async_connection
|
|
for options in RecordingAsyncPool.created
|
|
)
|
|
|
|
|
|
def test_synchronous_project_pools_check_connections_before_borrow(monkeypatch) -> None:
|
|
RecordingPool.created = []
|
|
monkeypatch.setattr(connection, "ConnectionPool", RecordingPool)
|
|
monkeypatch.setattr(connection, "_pools", OrderedDict())
|
|
monkeypatch.setattr(connection, "_pool_conninfo", {})
|
|
monkeypatch.setattr(connection, "_pool_borrows", {})
|
|
monkeypatch.setattr(connection, "_admin_pools", OrderedDict())
|
|
monkeypatch.setattr(connection, "_admin_pool_borrows", {})
|
|
monkeypatch.setattr(
|
|
connection,
|
|
"get_project_pgconn_string",
|
|
lambda db_name: f"postgresql://user:password@db.example/{db_name}",
|
|
)
|
|
|
|
connection.get_project_pool("tjwater_next")
|
|
connection.get_admin_pool()
|
|
|
|
monkeypatch.setattr(sync_pool, "ConnectionPool", RecordingPool)
|
|
monkeypatch.setattr(sync_pool, "_pools", OrderedDict())
|
|
monkeypatch.setattr(sync_pool, "_pool_conninfo", {})
|
|
monkeypatch.setattr(sync_pool, "_pool_borrows", {})
|
|
monkeypatch.setattr(
|
|
sync_pool,
|
|
"get_project_timescale_pgconn_string",
|
|
lambda db_name: f"postgresql://user:password@db.example/{db_name}",
|
|
)
|
|
sync_pool.get_timescale_pool("tjwater_next")
|
|
|
|
assert len(RecordingPool.created) == 3
|
|
assert [options["check"] for options in RecordingPool.created] == [
|
|
connection._check_connection,
|
|
connection._check_connection,
|
|
sync_pool._check_connection,
|
|
]
|