from contextlib import contextmanager import pytest from app.native.wndb.core import connection class _FakePool: def __init__(self, *, conninfo, **_kwargs): self.conninfo = conninfo self.closed = False self.borrowed = 0 self.returned = 0 def close(self): self.closed = True @contextmanager def connection(self): self.borrowed += 1 try: yield _FakeConnection() finally: self.returned += 1 class _FakeConnection: def __init__(self): self.transactions = 0 @contextmanager def transaction(self): self.transactions += 1 yield @pytest.fixture(autouse=True) def clear_pools(): connection._pools.clear() connection._pool_conninfo.clear() connection._pool_borrows.clear() connection._admin_pools.clear() connection._admin_pool_borrows.clear() yield connection._pools.clear() connection._pool_conninfo.clear() connection._pool_borrows.clear() connection._admin_pools.clear() connection._admin_pool_borrows.clear() def test_project_pool_is_reused_for_same_routed_dsn(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr(connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}") first = connection.get_project_pool("fengyang") second = connection.get_project_pool("fengyang") assert first is second assert first.conninfo == "dbname=fengyang" def test_project_pool_rebuilds_when_routed_dsn_changes(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) dsn = {"value": "host=old dbname=fengyang"} monkeypatch.setattr(connection, "get_project_pgconn_string", lambda *, db_name: dsn["value"]) old = connection.get_project_pool("fengyang") dsn["value"] = "host=new dbname=fengyang" new = connection.get_project_pool("fengyang") assert old.closed is True assert new is not old assert new.conninfo == "host=new dbname=fengyang" def test_project_connection_returns_connection_to_pool(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr(connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}") pool = connection.get_project_pool("fengyang") with connection.project_connection("fengyang"): assert pool.borrowed == 1 assert pool.returned == 0 assert pool.returned == 1 def test_project_transaction_reuses_one_pooled_connection(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}", ) pool = connection.get_project_pool("fengyang") with connection.project_transaction("fengyang") as transaction_conn: with connection.project_connection("fengyang") as nested_conn: assert nested_conn is transaction_conn assert transaction_conn.transactions == 1 assert pool.borrowed == 1 assert pool.returned == 1 def test_close_project_pool_removes_and_closes_pool(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr(connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}") pool = connection.get_project_pool("fengyang") connection.close_project_pool("fengyang") assert pool.closed is True assert "fengyang" not in connection._pools def test_admin_connection_is_pooled(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}", ) first = connection.get_admin_pool() second = connection.get_admin_pool() with connection.admin_connection(): pass assert first is second assert first.conninfo == "dbname=postgres" assert first.borrowed == 1 assert first.returned == 1 def test_close_all_closes_project_and_admin_pools(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}", ) project_pool = connection.get_project_pool("fengyang") admin_pool = connection.get_admin_pool() connection.close_all_project_pools() assert project_pool.closed is True assert admin_pool.closed is True assert connection._pools == {} assert connection._admin_pools == {} def test_admin_pools_are_isolated_by_routed_host(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) route = {"host": "one"} monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"host={route['host']} dbname={db_name}", ) first = connection.get_admin_pool() route["host"] = "two" second = connection.get_admin_pool() assert first is not second assert first.closed is False assert second.closed is False def test_project_pool_cache_evicts_least_recently_used_idle_pool(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr(connection.settings, "PROJECT_PG_CACHE_SIZE", 2) monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}", ) first = connection.get_project_pool("first") second = connection.get_project_pool("second") connection.get_project_pool("first") third = connection.get_project_pool("third") assert list(connection._pools) == ["first", "third"] assert second.closed is True assert first.closed is False assert third.closed is False def test_project_pool_cache_does_not_evict_active_pool(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) monkeypatch.setattr(connection.settings, "PROJECT_PG_CACHE_SIZE", 1) monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: f"dbname={db_name}", ) with connection.project_connection("active"): active = connection._pools["active"] connection.get_project_pool("new") assert active.closed is False assert set(connection._pools) == {"active", "new"} assert list(connection._pools) == ["new"] assert active.closed is True def test_route_health_check_does_not_close_active_pool(monkeypatch): monkeypatch.setattr(connection, "ConnectionPool", _FakePool) route = {"value": "host=old dbname=project"} monkeypatch.setattr( connection, "get_project_pgconn_string", lambda *, db_name: route["value"], ) with connection.project_connection("project"): pool = connection._pools["project"] route["value"] = "host=new dbname=project" assert connection.is_project_pool_open("project") is False assert pool.closed is False replacement = connection.get_project_pool("project") assert pool.closed is True assert replacement is not pool