TJWaterServerBinary/tests/unit/test_wndb_connection.py

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