refactor(db)!: finalize pooled WNDB v2 migration

tests/unit/test_pool_connection_checks.py

@@ -0,0 +1,104 @@
+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,
+    ]