Files
TJWaterServerBinary/tests/unit/test_wndb_connection.py
T
jiang b977bf6725
Server CI/CD / docker-image (push) Successful in 23s
Server CI/CD / deploy-fallback-log (push) Has been skipped
fix(db): validate cached project connections
2026-07-21 11:26:21 +08:00

122 lines
3.6 KiB
Python

import pytest
from app.native.wndb import connection
from app.native.wndb import database
from app.native.wndb import project
class _FakeCursor:
def __init__(self, connection):
self.connection = connection
def __enter__(self):
return self
def __exit__(self, exc_type, exc, tb):
return False
def execute(self, sql):
self.connection.executed.append(sql)
if self.connection.fail_ping and sql == "SELECT 1":
raise connection.pg.OperationalError("server closed the connection")
def fetchall(self):
return self.connection.rows
class _FakeConnection:
def __init__(self, rows=None, *, closed=False, fail_ping=False):
self.rows = list(rows or [])
self.closed = closed
self.fail_ping = fail_ping
self.executed = []
self.close_calls = 0
def cursor(self, row_factory=None):
if self.closed:
raise RuntimeError("the connection is closed")
return _FakeCursor(self)
def close(self):
self.close_calls += 1
self.closed = True
@pytest.fixture(autouse=True)
def clear_native_connections():
connection.g_conn_dict.clear()
connection._project_locks.clear()
yield
connection.g_conn_dict.clear()
connection._project_locks.clear()
def test_is_project_open_drops_closed_cached_connection():
connection.g_conn_dict["fengyang"] = _FakeConnection(closed=True)
assert project.is_project_open("fengyang") is False
assert "fengyang" not in connection.g_conn_dict
def test_open_connection_reuses_healthy_cached_connection(monkeypatch):
cached = _FakeConnection()
connection.g_conn_dict["fengyang"] = cached
def fail_connect(*, conninfo, autocommit):
raise AssertionError("cached connection should be reused")
monkeypatch.setattr(connection.pg, "connect", fail_connect)
assert connection.open_connection("fengyang") is cached
assert cached.executed == ["SELECT 1"]
def test_read_all_reopens_closed_cached_connection(monkeypatch):
stale = _FakeConnection(closed=True)
fresh = _FakeConnection(rows=[{"key": "DURATION", "value": "01:00:00"}])
connection.g_conn_dict["fengyang"] = stale
opened = []
def fake_connect(*, conninfo, autocommit):
opened.append((conninfo, autocommit))
return fresh
monkeypatch.setattr(connection.pg, "connect", fake_connect)
monkeypatch.setattr(
connection, "get_pgconn_string", lambda db_name: f"dbname={db_name}"
)
rows = database.read_all("fengyang", "select * from times")
assert rows == [{"key": "DURATION", "value": "01:00:00"}]
assert opened == [("dbname=fengyang", True)]
assert connection.g_conn_dict["fengyang"] is fresh
assert fresh.executed == ["select * from times"]
def test_read_all_reopens_cached_connection_when_health_check_fails(monkeypatch):
stale = _FakeConnection(fail_ping=True)
fresh = _FakeConnection(rows=[{"scheme_name": "base"}])
connection.g_conn_dict["fengyang"] = stale
opened = []
def fake_connect(*, conninfo, autocommit):
opened.append((conninfo, autocommit))
return fresh
monkeypatch.setattr(connection.pg, "connect", fake_connect)
monkeypatch.setattr(
connection, "get_pgconn_string", lambda db_name: f"dbname={db_name}"
)
rows = database.read_all("fengyang", "select * from scheme_list")
assert rows == [{"scheme_name": "base"}]
assert stale.executed == ["SELECT 1"]
assert stale.close_calls == 1
assert opened == [("dbname=fengyang", True)]
assert connection.g_conn_dict["fengyang"] is fresh
assert fresh.executed == ["select * from scheme_list"]