OrgTJWater/TJWaterServerBinary
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(db)!: finalize pooled WNDB v2 migration

Browse Source
This commit is contained in:
jiang
2026-08-27 17:26:22 +08:00
parent fa188af0b1
commit b74799a39d
105 changed files with 4988 additions and 5565 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/api/test_meta_endpoints.py
+8 -6
View File
@@ -2,7 +2,7 @@ from types import SimpleNamespace
from uuid import uuid4
from fastapi.testclient import TestClient
from sqlalchemy.exc import SQLAlchemyError
import psycopg
from tests.conftest import build_test_app, install_stub, load_module_from_path
@@ -15,7 +15,7 @@ def _load_meta_module(monkeypatch):
{
"ProjectContext": object,
"get_project_context": lambda: None,
"get_project_pg_session": lambda: None,
"get_project_pg_connection": lambda: None,
"get_project_timescale_connection": lambda: None,
"get_metadata_repository": lambda: None,
},
@@ -68,16 +68,18 @@ def test_meta_project_returns_map_extent(monkeypatch):
def test_meta_db_health_returns_503_for_postgres_errors(monkeypatch):
module = _load_meta_module(monkeypatch)
class BrokenSession:
async def execute(self, _query):
raise SQLAlchemyError("pg unavailable")
class BrokenConnection:
def cursor(self):
raise psycopg.OperationalError("pg unavailable")
class DummyTimescaleConnection:
def cursor(self):
raise AssertionError("timescale should not be queried after postgres failure")
app = build_test_app(module.router, "/api/v1")
app.dependency_overrides[module.get_project_pg_session] = lambda: BrokenSession()
app.dependency_overrides[module.get_project_pg_connection] = (
lambda: BrokenConnection()
)
app.dependency_overrides[module.get_project_timescale_connection] = lambda: DummyTimescaleConnection()
client = TestClient(app)
tests/api/test_model_import_endpoints.py
+47
View File
@@ -1,7 +1,9 @@
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
@@ -12,6 +14,7 @@ from app.auth.metadata_dependencies import (
)
from app.infra.db.metadb.repositories.metadata_repository import ProjectDbRouting
from app.infra.db.project_routing import get_project_pgconn_string
from app.native.wndb.core.database import MaterializedViewRefreshAfterCommitError
from tests.conftest import build_test_app
@@ -144,3 +147,47 @@ def test_model_update_uses_project_business_routing(monkeypatch):
"dsn": "postgresql://user:password@biz.example/routed_business",
}
repo.get_project_db_routing.assert_awaited_once_with(project_id, "biz_data")
def test_gb18030_upload_is_normalized_to_utf8() -> None:
content = "[TITLE]\n天津供水\n[JUNCTIONS]\n".encode("gb18030")
class FakeUpload:
filename = "model.inp"
async def read(self, _limit: int) -> bytes:
return content
normalized, filename = asyncio.run(model_import._read_upload(FakeUpload()))
assert filename == "model.inp"
assert normalized.decode("utf-8") == "[TITLE]\n天津供水\n[JUNCTIONS]\n"
def test_model_update_runs_blocking_import_in_threadpool(monkeypatch) -> None:
calls: list[tuple[object, tuple[object, ...]]] = []
async def fake_threadpool(function, *args):
calls.append((function, args))
monkeypatch.setattr(model_import, "run_in_threadpool", fake_threadpool)
asyncio.run(model_import._update_from_inp(b"[TITLE]\n", "demo"))
assert calls == [(model_import._update_from_inp_sync, (b"[TITLE]\n", "demo"))]
def test_committed_refresh_failure_is_not_wrapped_as_retryable_500(
monkeypatch,
) -> None:
error = MaterializedViewRefreshAfterCommitError("demo")
async def fail_update(_content: bytes, _project_code: str) -> None:
raise error
monkeypatch.setattr(model_import, "_update_from_inp", fail_update)
with pytest.raises(MaterializedViewRefreshAfterCommitError) as exc_info:
asyncio.run(model_import._apply_model_update(b"[TITLE]\n", "demo"))
assert exc_info.value is error
tests/api/test_openapi_contract.py
+43 -11
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import inspect
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4
@@ -12,6 +11,7 @@ from fastapi.testclient import TestClient
from app.api.v1.endpoints import simulation as simulation_endpoint
from app.api.pagination import PaginatedList
from app.api.problem_details import install_problem_details_handlers
from app.api.v1.rest_router import api_router, build_rest_router
from app.api.v1.router import api_router as source_api_router
from app.auth.project_dependencies import (
@@ -207,19 +207,51 @@ def test_valve_isolation_route_uses_the_isolation_handler() -> None:
assert route.name == "valve_isolation_endpoint"
def test_open_project_route_requires_business_and_timescale_routing() -> None:
route = next(
route
def test_legacy_project_pool_lifecycle_routes_are_not_published() -> None:
operations = {
(method, route.path)
for route in api_router.routes
if isinstance(route, APIRoute)
and route.path == "/projects/current"
and route.methods == {"POST"}
)
routing_parameter = inspect.signature(route.endpoint).parameters[
"_rest_project_routing"
]
for method in route.methods or set()
}
assert routing_parameter.default.dependency is get_project_simulation_routing
assert ("POST", "/projects/current") not in operations
assert ("DELETE", "/projects/current") not in operations
assert ("GET", "/projects/current/status") not in operations
def test_legacy_server_filesystem_inp_routes_are_not_published() -> None:
operations = {
(method, route.path)
for route in source_api_router.routes
if isinstance(route, APIRoute)
for method in route.methods or set()
}
assert operations.isdisjoint(
{
("POST", "/projects/current/imports"),
("POST", "/projects/current/exports/inp"),
("GET", "/projects/current/files/inp"),
}
)
def test_committed_write_refresh_failure_has_explicit_http_contract() -> None:
from app.native.wndb.core.database import MaterializedViewRefreshAfterCommitError
app = FastAPI()
install_problem_details_handlers(app)
@app.get("/probe")
def probe():
raise MaterializedViewRefreshAfterCommitError("project_a")
response = TestClient(app, raise_server_exceptions=False).get("/probe")
assert response.status_code == 503
assert response.headers["X-TJWater-Changes-Committed"] == "true"
assert response.json()["code"] == "materialized_view_refresh_failed_after_commit"
def test_valve_isolation_runtime_accepts_frontend_query(monkeypatch) -> None:
tests/api/test_project_endpoints.py
+17 -33
View File
@@ -18,6 +18,9 @@ class DummyChangeSet:
def _load_project_module(monkeypatch):
class DummyProjectContext:
project_code = "demo"
install_stub(monkeypatch, "app.services", package=True)
install_stub(
monkeypatch,
@@ -33,14 +36,9 @@ def _load_project_module(monkeypatch):
"have_project": lambda network: network == "demo",
"create_project": lambda network: None,
"delete_project": lambda network: None,
"is_project_open": lambda network: False,
"open_project": lambda network: None,
"close_project": lambda network: None,
"copy_project": lambda source, target: None,
"import_inp": lambda network, cs: {"ok": True},
"export_inp": lambda network, version: DummyChangeSet({"kind": "export"}),
"read_inp": lambda network, inp: True,
"dump_inp": lambda network, inp: True,
"get_all_vertices": lambda network: [],
"get_all_scada_info": lambda network: [],
"get_all_district_metering_areas": lambda network: [],
@@ -52,7 +50,12 @@ def _load_project_module(monkeypatch):
install_stub(
monkeypatch,
"app.auth.project_dependencies",
{"get_metadata_repository": lambda: None},
{
"ProjectContext": DummyProjectContext,
"get_metadata_repository": lambda: None,
"get_project_context": lambda: DummyProjectContext(),
"use_project_business_routing": lambda: None,
},
)
return load_module_from_path(
"tests_project_endpoints_module",
@@ -98,40 +101,21 @@ def test_project_info_returns_project_workspace(monkeypatch):
assert "geoserver" not in payload
def test_open_project_uses_unified_wndb_connection_path(monkeypatch):
def test_legacy_open_project_endpoint_is_removed(monkeypatch):
module = _load_project_module(monkeypatch)
called = []
monkeypatch.setattr(module, "open_project", lambda network: called.append(network))
client = TestClient(build_test_app(module.router, "/api/v1"))
response = client.post("/api/v1/projects/current", params={"network": "demo"})
assert response.status_code == 200
assert response.json() == "demo"
assert called == ["demo"]
assert response.status_code == 405
def test_project_lock_lifecycle(monkeypatch):
def test_legacy_physical_project_routes_are_removed(monkeypatch):
module = _load_project_module(monkeypatch)
module.lockedPrjs.clear()
client = TestClient(build_test_app(module.router, "/api/v1"))
first_lock = client.post("/api/v1/projects/current/lock", params={"network": "demo"})
second_lock = client.post("/api/v1/projects/current/lock", params={"network": "demo"})
locked_by_me = client.get(
"/api/v1/projects/current/lock/ownership",
params={"network": "demo"},
)
unlock = client.delete(
"/api/v1/projects/current/lock",
params={"network": "demo"},
)
locked = client.get("/api/v1/projects/current/lock", params={"network": "demo"})
assert first_lock.json() == 0
assert second_lock.json() == 1
assert locked_by_me.json() is True
assert unlock.json() is True
assert locked.json() is False
assert client.get("/api/v1/project-codes").status_code == 404
assert client.get("/api/v1/projects/existence").status_code == 404
assert client.post("/api/v1/project-copies").status_code == 404
assert client.get("/api/v1/projects/current/lock").status_code == 404
assert client.post("/api/v1/project-conversions").status_code == 404
tests/api/test_wndb_endpoint_threading.py
+68
View File
@@ -0,0 +1,68 @@
from threading import get_ident
from uuid import uuid4
from fastapi import FastAPI
from fastapi.testclient import TestClient
from app.api.v1.endpoints.components import curves
from app.api.v1.rest_router import api_router
from app.auth.project_dependencies import (
ProjectContext,
get_project_business_routing,
get_project_context,
)
from app.infra.db.project_routing import ActiveProjectRouting, get_active_project_routing
def test_sync_wndb_endpoint_parses_body_and_runs_in_worker_thread(monkeypatch) -> None:
call: dict[str, object] = {}
project_id = uuid4()
user_id = uuid4()
context = ProjectContext(
project_id=project_id,
project_code="project_a",
user_id=user_id,
project_role="member",
)
routing = ActiveProjectRouting(
project_code="project_a",
business_dsn="postgresql://user:password@db.example/project_a",
)
async def override_context() -> ProjectContext:
call["event_loop_thread"] = get_ident()
return context
async def override_routing() -> ActiveProjectRouting:
return routing
def fake_add_curve(network, changes):
call["worker_thread"] = get_ident()
call["network"] = network
call["operations"] = changes.operations
call["routing"] = get_active_project_routing()
return {"ok": True}
monkeypatch.setattr(curves, "add_curve", fake_add_curve)
app = FastAPI()
app.include_router(api_router)
app.dependency_overrides[get_project_context] = override_context
app.dependency_overrides[get_project_business_routing] = override_routing
with TestClient(app) as client:
response = client.post(
"/curves",
params={"curve": "C-1"},
json={"points": [[0, 10], [1, 20]]},
)
assert response.status_code == 201
assert response.json() == {"ok": True}
assert call == {
"event_loop_thread": call["event_loop_thread"],
"worker_thread": call["worker_thread"],
"network": "project_a",
"operations": [{"id": "C-1", "points": [[0, 10], [1, 20]]}],
"routing": routing,
}
assert call["worker_thread"] != call["event_loop_thread"]
tests/integration/test_database_pooling_live.py
+121 -3
View File
@@ -1,12 +1,17 @@
import asyncio
import os
from concurrent.futures import ThreadPoolExecutor
from uuid import uuid4
import pytest
from psycopg import connect
from app.core.config import get_pgconn_string
from app.infra.db.dynamic_manager import ProjectConnectionManager
from app.infra.db.timescaledb.sync_pool import timescale_connection
from app.native.wndb.commands.api import delete_pattern_cascade
from app.native.wndb.core.connection import project_connection, project_transaction
from app.native.wndb.core.projects import have_project, temporary_project_database
from app.native.wndb.core.database import ChangeSet, g_delete_prefix, write
from app.native.wndb.model import demands, junctions, patterns
from app.services.scheme_management import create_analysis_run, update_analysis_run
@@ -17,7 +22,7 @@ pytestmark = pytest.mark.skipif(
reason="set RUN_DB_INTEGRATION=1 to test configured PostgreSQL databases",
)
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_next")
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_v2")
def _read_business_database(_: int) -> str:
@@ -46,6 +51,60 @@ def test_timeseries_pool_handles_concurrent_borrows() -> None:
assert names == [PROJECT] * 64
def test_temporary_project_clone_copies_model_scada_and_views() -> None:
def counts(project: str) -> dict:
with project_connection(project) as conn, conn.cursor() as cur:
cur.execute(
"""
select
(select count(*) from network.nodes) as nodes,
(select count(*) from network.links) as links,
(select count(*) from asset.scada_devices) as scada,
(select count(*) from gis.junctions) as mv_junctions,
(select count(*) from gis.pipes) as mv_pipes
"""
)
return dict(cur.fetchone())
source_counts = counts(PROJECT)
temporary = None
with temporary_project_database(PROJECT, "clone_validation") as temporary:
assert counts(temporary) == source_counts
assert temporary is not None
assert have_project(temporary) is False
def test_dynamic_pool_replaces_terminated_idle_connection() -> None:
async def exercise_pool() -> None:
manager = ProjectConnectionManager()
dsn = get_pgconn_string(db_name=PROJECT)
try:
project_id = uuid4()
async with manager.pg_connection(
project_id, "biz_data", dsn, 1, 1
) as conn:
async with conn.cursor() as cur:
await cur.execute("select pg_backend_pid()")
backend_pid = int((await cur.fetchone())["pg_backend_pid"])
with connect(dsn, autocommit=True) as admin_conn:
with admin_conn.cursor() as cur:
cur.execute("select pg_terminate_backend(%s)", (backend_pid,))
assert cur.fetchone()[0] is True
async with manager.pg_connection(
project_id, "biz_data", dsn, 1, 1
) as conn:
async with conn.cursor() as cur:
await cur.execute("select current_database()")
assert (await cur.fetchone())["current_database"] == PROJECT
finally:
await manager.close_all()
asyncio.run(exercise_pool())
def test_nested_wndb_writes_roll_back_as_one_transaction() -> None:
with pytest.raises(RuntimeError, match="force rollback"):
with project_transaction(PROJECT) as conn:
@@ -186,10 +245,55 @@ def test_wndb_ordered_detail_tables_use_parent_scoped_primary_keys() -> None:
assert actual == expected
def test_gis_unified_views_cover_all_materialized_network_layers() -> None:
with project_connection(PROJECT) as conn, conn.cursor() as cur:
cur.execute(
"""
SELECT
(SELECT COUNT(*) FROM gis.network_nodes) AS nodes,
(SELECT COUNT(*) FROM gis.junctions)
+ (SELECT COUNT(*) FROM gis.reservoirs)
+ (SELECT COUNT(*) FROM gis.tanks) AS source_nodes,
(SELECT COUNT(*) FROM gis.network_links) AS links,
(SELECT COUNT(*) FROM gis.pipes)
+ (SELECT COUNT(*) FROM gis.pumps)
+ (SELECT COUNT(*) FROM gis.valves) AS source_links
"""
)
counts = cur.fetchone()
cur.execute(
"""
SELECT obj_description('gis.network_nodes'::regclass) AS node_comment,
obj_description('gis.network_links'::regclass) AS link_comment
"""
)
comments = cur.fetchone()
cur.execute(
"""
SELECT COUNT(*) AS undocumented_columns
FROM pg_attribute
WHERE attrelid = ANY(
ARRAY['gis.network_nodes'::regclass, 'gis.network_links'::regclass]
)
AND attnum > 0
AND NOT attisdropped
AND col_description(attrelid, attnum) IS NULL
"""
)
undocumented_columns = cur.fetchone()["undocumented_columns"]
assert counts["nodes"] == counts["source_nodes"]
assert counts["links"] == counts["source_links"]
assert comments["node_comment"]
assert comments["link_comment"]
assert undocumented_columns == 0
def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
suffix = uuid4()
junction_id = f"integration-junction-{suffix}"
pattern_id = f"integration-cascade-{suffix}"
retained_pattern_id = f"integration-retained-{suffix}"
with pytest.raises(RuntimeError, match="force rollback"):
with project_transaction(PROJECT):
@@ -203,6 +307,10 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
PROJECT,
ChangeSet({"id": pattern_id, "factors": [1.0]}),
)
patterns.add_pattern(
PROJECT,
ChangeSet({"id": retained_pattern_id, "factors": [1.0]}),
)
demands.set_demand(
PROJECT,
ChangeSet(
@@ -213,7 +321,12 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
"demand": 1.0,
"pattern": pattern_id,
"category": "integration",
}
},
{
"demand": 2.0,
"pattern": retained_pattern_id,
"category": "retained",
},
],
}
),
@@ -226,7 +339,12 @@ def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
assert patterns.get_pattern(PROJECT, pattern_id) == {}
assert demands.get_demand(PROJECT, junction_id)["demands"] == [
{"demand": 1.0, "pattern": None, "category": "integration"}
{"demand": 1.0, "pattern": None, "category": "integration"},
{
"demand": 2.0,
"pattern": retained_pattern_id,
"category": "retained",
},
]
assert result.operations[-1] == {
"operation": "delete",
tests/unit/test_age_analysis.py
+33 -10
View File
@@ -1,4 +1,5 @@
import json
from contextlib import contextmanager
from tests.conftest import install_stub, load_module_from_path
@@ -29,7 +30,6 @@ def _load_scenarios_module(monkeypatch):
"SOURCE_TYPE_SETPOINT": "SOURCE_TYPE_SETPOINT",
"add_pattern": lambda *args, **kwargs: None,
"add_source": lambda *args, **kwargs: None,
"close_project": lambda *args, **kwargs: None,
"copy_project": lambda *args, **kwargs: None,
"delete_project": lambda *args, **kwargs: None,
"get_demand": lambda *args, **kwargs: None,
@@ -42,8 +42,6 @@ def _load_scenarios_module(monkeypatch):
"get_time": lambda *args, **kwargs: None,
"have_project": lambda *args, **kwargs: False,
"is_junction": lambda *args, **kwargs: False,
"is_project_open": lambda *args, **kwargs: False,
"open_project": lambda *args, **kwargs: None,
"set_demand": lambda *args, **kwargs: None,
"set_emitter": lambda *args, **kwargs: None,
"set_option": lambda *args, **kwargs: None,
@@ -61,12 +59,11 @@ def test_age_analysis_passes_duration_by_keyword(monkeypatch):
module = _load_scenarios_module(monkeypatch)
captured = {}
monkeypatch.setattr(module, "copy_project", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "open_project", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "close_project", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "delete_project", lambda *args, **kwargs: None)
monkeypatch.setattr(module, "have_project", lambda *args, **kwargs: False)
monkeypatch.setattr(module, "is_project_open", lambda *args, **kwargs: False)
@contextmanager
def fake_temporary_project(project, purpose):
yield f"{purpose}_{project}_run"
monkeypatch.setattr(module, "temporary_project_database", fake_temporary_project)
def fake_run_simulation_ex(*args, **kwargs):
captured["args"] = args
@@ -78,7 +75,7 @@ def test_age_analysis_passes_duration_by_keyword(monkeypatch):
module.age_analysis("demo", "2026-06-03T07:00:00+08:00", 300)
assert captured["args"] == (
"age_Anal_demo",
"age_analysis_demo_run",
"realtime",
"2026-06-03T07:00:00+08:00",
)
@@ -86,3 +83,29 @@ def test_age_analysis_passes_duration_by_keyword(monkeypatch):
"duration": 300,
"downloading_prohibition": True,
}
def test_isolated_analysis_cleans_database_after_early_return(monkeypatch):
module = _load_scenarios_module(monkeypatch)
lifecycle: list[tuple[str, str]] = []
@contextmanager
def fake_temporary_project(project, purpose):
lifecycle.append(("create", purpose))
try:
yield "isolated_run"
finally:
lifecycle.append(("delete", purpose))
monkeypatch.setattr(
module, "temporary_project_database", fake_temporary_project
)
@module._isolated_analysis("probe")
def return_early(name, *, _temporary_project=None):
assert name == "demo"
assert _temporary_project == "isolated_run"
return "done"
assert return_early("demo") == "done"
assert lifecycle == [("create", "probe"), ("delete", "probe")]
tests/unit/test_analysis_simulation.py
+90 -2
View File
@@ -1,7 +1,11 @@
import inspect
import json
from contextlib import contextmanager
from unittest.mock import Mock
from uuid import uuid4
import pytest
def test_run_simulation_exposes_explicit_valve_control():
from app.services import simulation
@@ -9,6 +13,34 @@ def test_run_simulation_exposes_explicit_valve_control():
assert "valve_control" in inspect.signature(simulation.run_simulation).parameters
def test_extended_runner_cleans_temporary_database_after_failure(monkeypatch):
from app.algorithms.simulation import runner
lifecycle: list[tuple[str, str]] = []
@contextmanager
def temporary_project(project: str, purpose: str):
lifecycle.append(("create", project))
try:
yield "isolated_project"
finally:
lifecycle.append(("delete", project))
monkeypatch.setattr(runner, "temporary_project_database", temporary_project)
@runner._clean_extended_simulation
def fail(name, simulation_type, *, _temporary_project=None):
assert name == "demo"
assert simulation_type == "extended"
assert _temporary_project == "isolated_project"
raise RuntimeError("simulation failed")
with pytest.raises(RuntimeError, match="simulation failed"):
fail("demo", "extended")
assert lifecycle == [("create", "demo"), ("delete", "demo")]
def test_apply_valve_control_matches_runner_semantics(monkeypatch):
from app.services import simulation
@@ -46,12 +78,58 @@ def test_apply_valve_control_matches_runner_semantics(monkeypatch):
assert updates["V-k"]["setting"] == 0.1036 * pow(0.5, -3.105)
def test_primary_demand_update_preserves_additional_categories():
from app.services import simulation
demand_set = {
"junction": "J1",
"demands": [
{"demand": 1.0, "pattern": "P1", "category": "domestic"},
{"demand": 2.0, "pattern": "P2", "category": "industrial"},
],
}
simulation._primary_demand(demand_set)["demand"] = 3.0
assert demand_set["demands"] == [
{"demand": 3.0, "pattern": "P1", "category": "domestic"},
{"demand": 2.0, "pattern": "P2", "category": "industrial"},
]
assert simulation._primary_demand_pattern(demand_set) == "P1"
def test_primary_demand_is_created_for_empty_junction():
from app.services import simulation
demand_set = {"junction": "J1", "demands": []}
primary = simulation._primary_demand(demand_set)
assert primary == {"demand": 0.0, "pattern": None, "category": None}
with pytest.raises(ValueError, match="has no demand pattern"):
simulation._primary_demand_pattern(demand_set)
def test_extended_simulation_stores_results_by_run_id(monkeypatch):
from app.services import simulation
run_id = uuid4()
storage_calls: list[tuple] = []
monkeypatch.setattr(simulation, "open_project", lambda name: None)
transaction_calls: list[tuple[str, str]] = []
@contextmanager
def project_transaction(name):
transaction_calls.append(("begin", name))
try:
yield object()
finally:
transaction_calls.append(("end", name))
refresh_mock = Mock()
monkeypatch.setattr(simulation, "project_transaction", project_transaction)
monkeypatch.setattr(
simulation, "refresh_materialized_views_after_commit", refresh_mock
)
monkeypatch.setattr(
simulation,
"get_time",
@@ -104,6 +182,8 @@ def test_extended_simulation_stores_results_by_run_id(monkeypatch):
assert kwargs["db_name"] == "demo"
assert returned_run_id == run_id
assert lifecycle_calls[-1][1]["status"] == "completed"
assert transaction_calls == [("begin", "demo"), ("end", "demo")]
refresh_mock.assert_called_once_with("demo")
def test_extended_simulation_marks_run_failed_when_result_storage_fails(monkeypatch):
@@ -111,7 +191,15 @@ def test_extended_simulation_marks_run_failed_when_result_storage_fails(monkeypa
run_id = uuid4()
lifecycle_calls: list[tuple] = []
monkeypatch.setattr(simulation, "open_project", lambda name: None)
@contextmanager
def project_transaction(_name):
yield object()
monkeypatch.setattr(simulation, "project_transaction", project_transaction)
monkeypatch.setattr(
simulation, "refresh_materialized_views_after_commit", lambda _name: None
)
monkeypatch.setattr(
simulation,
"get_time",
tests/unit/test_clean_projects_cli.py
+17
View File
@@ -0,0 +1,17 @@
import pytest
from scripts.clean_projects import parse_args
def test_clean_projects_requires_explicit_confirmation() -> None:
with pytest.raises(SystemExit) as exc_info:
parse_args(["temporary_project"])
assert exc_info.value.code == 2
def test_clean_projects_accepts_exact_targets_after_confirmation() -> None:
args = parse_args(["--yes", "temp_a", "temp_b"])
assert args.yes is True
assert args.projects == ["temp_a", "temp_b"]
tests/unit/test_dynamic_manager.py
+98 -8
View File
@@ -1,12 +1,102 @@
from app.infra.db.dynamic_manager import ProjectConnectionManager
import asyncio
from contextlib import asynccontextmanager
from uuid import uuid4
import pytest
from app.infra.db import dynamic_manager
def test_normalize_pg_url_preserves_password():
manager = ProjectConnectionManager()
class FakeAsyncPool:
created: list["FakeAsyncPool"] = []
url = manager._normalize_pg_url(
"postgresql://tjwater:secret@192.168.1.114:5433/tjwater"
)
def __init__(self, **kwargs) -> None:
self.kwargs = kwargs
self.closed = False
self.created.append(self)
assert url == "postgresql+psycopg://tjwater:secret@192.168.1.114:5433/tjwater"
assert "***" not in url
async def open(self) -> None:
return None
async def close(self) -> None:
self.closed = True
@asynccontextmanager
async def connection(self):
yield object()
def test_active_project_pool_is_not_evicted(monkeypatch) -> None:
async def exercise() -> None:
manager = dynamic_manager.ProjectConnectionManager()
first_id = uuid4()
second_id = uuid4()
async with manager.pg_connection(first_id, "biz_data", "dsn-1", 1, 2):
first_pool = manager._pg_raw_cache[
dynamic_manager.CacheKey(first_id, "biz_data")
].pool
async with manager.pg_connection(
second_id, "biz_data", "dsn-2", 1, 2
):
assert first_pool.closed is False
assert len(manager._pg_raw_cache) == 2
assert first_pool.closed is False
assert list(manager._pg_raw_cache) == [
dynamic_manager.CacheKey(first_id, "biz_data")
]
await manager.close_all()
FakeAsyncPool.created = []
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
monkeypatch.setattr(dynamic_manager.settings, "PROJECT_PG_CACHE_SIZE", 1)
asyncio.run(exercise())
def test_active_pool_uses_generation_replacement(monkeypatch) -> None:
async def exercise() -> None:
manager = dynamic_manager.ProjectConnectionManager()
project_id = uuid4()
async with manager.pg_connection(
project_id, "biz_data", "old-dsn", 1, 2
):
old_pool = FakeAsyncPool.created[0]
async with manager.pg_connection(
project_id, "biz_data", "new-dsn", 1, 2
):
assert old_pool.closed is False
assert len(manager._retired_pg) == 1
assert old_pool.closed is True
assert manager._retired_pg == []
await manager.close_all()
FakeAsyncPool.created = []
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
asyncio.run(exercise())
def test_close_project_does_not_interrupt_active_borrow(monkeypatch) -> None:
async def exercise() -> None:
manager = dynamic_manager.ProjectConnectionManager()
project_id = uuid4()
key = dynamic_manager.CacheKey(project_id, "iot_data")
async with manager.timescale_connection(
project_id, "iot_data", "ts-dsn", 1, 2
):
pool = manager._ts_cache[key].pool
assert await manager.close_project(project_id) is False
assert pool.closed is False
assert await manager.close_project(project_id) is True
assert pool.closed is True
assert key not in manager._ts_cache
FakeAsyncPool.created = []
monkeypatch.setattr(dynamic_manager, "AsyncConnectionPool", FakeAsyncPool)
asyncio.run(exercise())
tests/unit/test_pool_connection_checks.py
+104
View File
@@ -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,
]
tests/unit/test_project_routing.py
+34 -3
View File
@@ -5,7 +5,9 @@ from app.infra.db.project_routing import (
ActiveProjectRouting,
activate_project_routing,
get_active_project_routing,
get_project_database_name,
get_project_pgconn_string,
get_project_template_database_name,
get_project_timescale_pgconn_string,
)
@@ -35,9 +37,16 @@ def test_project_database_uses_exact_routing_dsn_for_project_code() -> None:
)
def test_business_template_keeps_server_and_timescale_ignores_legacy_db_name() -> None:
def test_business_template_keeps_server_and_timescale_ignores_legacy_db_name(
monkeypatch,
) -> None:
monkeypatch.setattr(
"app.infra.db.project_routing.settings.WNDB_TEMPLATE_DB_NAME",
"tjwater_v2_template",
)
with activate_project_routing(_routing()):
business = conninfo_to_dict(get_project_pgconn_string("project_a_template"))
template_name = get_project_template_database_name("project_a")
business = conninfo_to_dict(get_project_pgconn_string(template_name))
timescale = conninfo_to_dict(
get_project_timescale_pgconn_string("temporary_scheme")
)
@@ -45,7 +54,7 @@ def test_business_template_keeps_server_and_timescale_ignores_legacy_db_name() -
assert business == {
"user": "biz_user",
"password": "biz_password",
"dbname": "project_a_template",
"dbname": "tjwater_v2_template",
"host": "biz.example",
"port": "5432",
"sslmode": "require",
@@ -60,6 +69,28 @@ def test_business_template_keeps_server_and_timescale_ignores_legacy_db_name() -
}
def test_project_code_resolves_to_physical_business_database(monkeypatch) -> None:
monkeypatch.setattr(
"app.infra.db.project_routing.settings.WNDB_TEMPLATE_DB_NAME",
"tjwater_v2_template",
)
with activate_project_routing(_routing()):
assert get_project_database_name("project_a") == "biz_database"
assert get_project_database_name("temporary_run") == "temporary_run"
assert get_project_template_database_name("project_a") == "tjwater_v2_template"
def test_template_falls_back_to_config_outside_project_routing(monkeypatch) -> None:
monkeypatch.setattr(
"app.infra.db.project_routing.settings.WNDB_TEMPLATE_DB_NAME",
"tjwater_v2_template",
)
assert get_project_template_database_name("ignored-project-code") == (
"tjwater_v2_template"
)
def test_project_routing_is_nested_and_request_local() -> None:
first = _routing("project_a")
second = _routing("project_b")
tests/unit/test_project_scada_metadata.py
+22 -6
View File
@@ -33,12 +33,19 @@ def _patch_project_scadas(monkeypatch):
def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
_patch_project_scadas(monkeypatch)
query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
query_mock = AsyncMock(
return_value={"J1": [{"time": START_TIME, "value": 26.5}]}
)
monkeypatch.setattr(
composite_queries.RealtimeRepository,
"get_node_field_by_time_range",
"get_node_fields_by_ids_time_range",
query_mock,
)
monkeypatch.setattr(
composite_queries.RealtimeRepository,
"get_link_fields_by_ids_time_range",
AsyncMock(return_value={}),
)
result = asyncio.run(
composite_queries.CompositeQueries.get_scada_associated_realtime_simulation_data(
@@ -55,17 +62,22 @@ def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
assert query_mock.await_args.args[1:] == (
START_TIME,
END_TIME,
"J1",
["J1"],
"pressure",
)
def test_analysis_scada_simulation_uses_current_project_metadata(monkeypatch):
_patch_project_scadas(monkeypatch)
query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
async def query_series(_conn, _run_id, element_type, element_ids, *_args):
if element_type == "node":
return {"J1": [{"time": START_TIME, "value": 26.5}]}
return {}
query_mock = AsyncMock(side_effect=query_series)
monkeypatch.setattr(
composite_queries.AnalysisResultsRepository,
"get_node_series",
"get_series_by_ids",
query_mock,
)
@@ -81,7 +93,11 @@ def test_analysis_scada_simulation_uses_current_project_metadata(monkeypatch):
)
assert result[PROJECT_SCADA["device_id"]][0]["scada_id"] == PROJECT_SCADA["device_id"]
assert query_mock.await_args.args[2:] == ("J1", START_TIME, END_TIME, "pressure")
assert query_mock.await_count == 2
node_call = next(
call for call in query_mock.await_args_list if call.args[2] == "node"
)
assert node_call.args[3:] == (["J1"], START_TIME, END_TIME, "pressure")
def test_element_scada_query_uses_current_project_metadata(monkeypatch):
tests/unit/test_realtime_repository.py
+74 -1
View File
@@ -1,5 +1,5 @@
import asyncio
from contextlib import contextmanager
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime, timezone
from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository
@@ -71,12 +71,23 @@ def test_get_nodes_by_time_range_normalizes_inputs_to_utc():
class _SyncTransactionConnection:
def __init__(self):
self.transactions = 0
self.calls: list[tuple[str, tuple]] = []
@contextmanager
def transaction(self):
self.transactions += 1
yield
@contextmanager
def cursor(self):
connection = self
class Cursor:
def execute(self, query, params):
connection.calls.append((query, params))
yield Cursor()
def test_realtime_node_and_link_replacement_share_outer_transaction(monkeypatch):
conn = _SyncTransactionConnection()
@@ -101,3 +112,65 @@ def test_realtime_node_and_link_replacement_share_outer_transaction(monkeypatch)
assert conn.transactions == 1
assert calls == ["nodes", "links"]
assert [query for query, _params in conn.calls] == [
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
"DELETE FROM realtime.node_results WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
]
def test_realtime_batch_rejects_multiple_timestamps_before_writing():
data = [
{"time": "2026-06-01T00:00:00Z", "id": "N1"},
{"time": "2026-06-01T00:15:00Z", "id": "N2"},
]
try:
RealtimeRepository.insert_nodes_batch_sync(object(), data)
except ValueError as exc:
assert str(exc) == "Realtime batch must contain exactly one timestamp"
else:
raise AssertionError("multiple realtime timestamps were accepted")
class _AsyncEmptySnapshotConnection:
def __init__(self):
self.calls: list[tuple[str, tuple]] = []
@asynccontextmanager
async def transaction(self):
yield
def cursor(self):
connection = self
class Cursor:
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return False
async def execute(self, query, params):
connection.calls.append((query, params))
return Cursor()
def test_empty_realtime_side_is_deleted_as_part_of_snapshot_replacement():
conn = _AsyncEmptySnapshotConnection()
asyncio.run(
RealtimeRepository.store_realtime_simulation_result(
conn,
node_result_list=[],
link_result_list=[],
result_start_time="2026-06-01T00:00:00Z",
)
)
assert [query for query, _params in conn.calls] == [
"SELECT pg_advisory_xact_lock(hashtextextended(%s::text, 1))",
"DELETE FROM realtime.node_results WHERE time = %s",
"DELETE FROM realtime.link_results WHERE time = %s",
]
tests/unit/test_scada_cleaning.py
+4 -3
View File
@@ -32,9 +32,10 @@ def test_clean_scada_uses_current_project_metadata(monkeypatch):
),
)
update_mock = AsyncMock()
update_mock.return_value = 1
monkeypatch.setattr(
composite_queries.ScadaRepository,
"update_scada_field",
"update_scada_field_batch",
update_mock,
)
monkeypatch.setattr(
@@ -111,7 +112,7 @@ def test_clean_scada_rejects_zero_database_updates(monkeypatch):
update_mock = AsyncMock()
monkeypatch.setattr(
composite_queries.ScadaRepository,
"update_scada_field",
"update_scada_field_batch",
update_mock,
)
monkeypatch.setattr(
@@ -157,7 +158,7 @@ def test_clean_scada_propagates_write_failures(monkeypatch):
)
monkeypatch.setattr(
composite_queries.ScadaRepository,
"update_scada_field",
"update_scada_field_batch",
AsyncMock(side_effect=RuntimeError("database write failed")),
)
monkeypatch.setattr(
tests/unit/test_scada_repository.py
+28
View File
@@ -86,3 +86,31 @@ def test_update_scada_field_skips_insert_when_update_succeeds():
assert len(conn.cursor_instance.calls) == 1
assert "UPDATE scada.measurements SET" in conn.cursor_instance.calls[0][0]
def test_update_scada_field_batch_uses_one_set_based_statement():
ScadaRepository = _load_scada_repository()
conn = _FakeConnection(initial_rowcount=2)
first_time = datetime(2026, 1, 1, 0, 0, tzinfo=timezone.utc)
second_time = datetime(2026, 1, 1, 0, 5, tzinfo=timezone.utc)
updated = asyncio.run(
ScadaRepository.update_scada_field_batch(
conn,
[
(first_time, "170490", 26.5),
(second_time, "170491", 27.0),
],
"cleaned_value",
)
)
assert updated == 2
assert len(conn.cursor_instance.calls) == 1
query, params = conn.cursor_instance.calls[0]
assert "unnest" in query.lower()
assert params == (
[first_time, second_time],
["170490", "170491"],
[26.5, 27.0],
)
tests/unit/test_wndb_batch_transactions.py
+212 -4
View File
@@ -3,7 +3,9 @@ from contextlib import contextmanager
import pytest
from app.native.wndb.commands import executor
from app.native.wndb.core import database
from app.native.wndb.core.database import ChangeSet
from app.native.wndb.model import junctions, pipes, pumps, reservoirs, tanks, valves
def test_batch_commits_before_materialized_view_refresh(monkeypatch) -> None:
@@ -15,7 +17,7 @@ def test_batch_commits_before_materialized_view_refresh(monkeypatch) -> None:
yield object()
events.append("transaction-exit")
monkeypatch.setattr(executor, "project_transaction", fake_transaction)
monkeypatch.setattr(executor, "model_mutation_transaction", fake_transaction)
monkeypatch.setattr(
executor,
"expand_command",
@@ -28,7 +30,7 @@ def test_batch_commits_before_materialized_view_refresh(monkeypatch) -> None:
)
monkeypatch.setattr(
executor,
"refresh_materialized_views",
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
@@ -56,7 +58,7 @@ def test_failed_batch_does_not_refresh_materialized_views(monkeypatch) -> None:
finally:
events.append("transaction-exit")
monkeypatch.setattr(executor, "project_transaction", fake_transaction)
monkeypatch.setattr(executor, "model_mutation_transaction", fake_transaction)
monkeypatch.setattr(
executor,
"expand_command",
@@ -70,7 +72,7 @@ def test_failed_batch_does_not_refresh_materialized_views(monkeypatch) -> None:
monkeypatch.setattr(executor, "_execute_update_command", fail_write)
monkeypatch.setattr(
executor,
"refresh_materialized_views",
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
@@ -81,3 +83,209 @@ def test_failed_batch_does_not_refresh_materialized_views(monkeypatch) -> None:
)
assert events == ["transaction-enter", "write", "transaction-exit"]
def test_batch_option_update_does_not_refresh_materialized_views(monkeypatch) -> None:
events: list[str] = []
@contextmanager
def fake_transaction(_name: str):
events.append("transaction-enter")
yield object()
events.append("transaction-exit")
monkeypatch.setattr(executor, "model_mutation_transaction", fake_transaction)
monkeypatch.setattr(executor, "expand_command", lambda _name, cs: cs)
monkeypatch.setattr(
executor,
"_execute_update_command",
lambda _name, _change_set: events.append("write") or ChangeSet(),
)
monkeypatch.setattr(
executor,
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
executor.execute_batch_commands(
"project_a",
ChangeSet({"operation": "update", "type": "option", "id": "duration"}),
)
assert events == ["transaction-enter", "write", "transaction-exit"]
def test_model_mutation_lock_is_acquired_once_per_transaction(monkeypatch) -> None:
state = {"held": False}
statements: list[tuple[str, tuple[str]]] = []
class FakeCursor:
def __enter__(self):
return self
def __exit__(self, *_args):
return None
def execute(self, statement: str, params: tuple[str]):
statements.append((statement, params))
class FakeConnection:
def cursor(self):
return FakeCursor()
monkeypatch.setattr(
database,
"is_model_mutation_lock_active",
lambda _name: state["held"],
raising=False,
)
monkeypatch.setattr(
database,
"mark_model_mutation_lock_active",
lambda _name: state.__setitem__("held", True),
raising=False,
)
monkeypatch.setattr(database, "get_project_database_name", lambda name: name)
conn = FakeConnection()
database.acquire_model_mutation_lock(conn, "project_a")
database.acquire_model_mutation_lock(conn, "project_a")
assert len(statements) == 1
def test_locked_command_builds_after_lock_and_refreshes_after_commit(
monkeypatch,
) -> None:
events: list[str] = []
@contextmanager
def fake_model_transaction(_name: str):
events.append("lock")
yield object()
events.append("commit")
def build_command() -> database.DatabaseCommand:
events.append("read-and-build")
return database.DatabaseCommand(
"UPDATE network.junctions SET elevation = 1",
[{"operation": "update", "type": "junction", "id": "J1"}],
)
monkeypatch.setattr(database, "model_mutation_transaction", fake_model_transaction)
monkeypatch.setattr(
database,
"is_project_transaction_active",
lambda _name: False,
)
monkeypatch.setattr(
database,
"execute_command",
lambda _name, command: events.append("write")
or ChangeSet.from_list(command.changes),
)
monkeypatch.setattr(
database,
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
result = database.execute_locked_command("project_a", build_command)
assert events == ["lock", "read-and-build", "write", "commit", "refresh"]
assert result.operations[0]["id"] == "J1"
def test_pipe_patch_reads_under_lock_and_updates_only_supplied_columns(
monkeypatch,
) -> None:
events: list[str] = []
current = {
"id": "P1",
"node1": "J1",
"node2": "J2",
"length": 10.0,
"diameter": 100.0,
"roughness": 120.0,
"minor_loss": 0.0,
"status": "OPEN",
}
def fake_locked_command(_name: str, builder):
events.append("lock")
command = builder()
assert command is not None
captured.append(command)
events.append("refresh")
return ChangeSet.from_list(command.changes)
def fake_get_pipe(_name: str, _id: str):
events.append("read")
return current.copy()
captured: list[database.DatabaseCommand] = []
monkeypatch.setattr(
pipes,
"execute_locked_command",
fake_locked_command,
)
monkeypatch.setattr(pipes, "get_pipe", fake_get_pipe)
result = pipes.set_pipe(
"project_a",
ChangeSet({"operation": "update", "type": "pipe", "id": "P1", "length": 20}),
)
assert events == ["lock", "read", "refresh"]
assert len(captured) == 1
assert "length = 20.0" in captured[0].sql
assert "diameter" not in captured[0].sql
assert "update network.links" not in captured[0].sql.lower()
assert result.operations[0]["length"] == 20.0
@pytest.mark.parametrize(
("module", "setter_name", "getter_name", "builder_name"),
[
(junctions, "set_junction", "get_junction", "_set_junction"),
(reservoirs, "set_reservoir", "get_reservoir", "_set_reservoir"),
(tanks, "set_tank", "get_tank", "_set_tank"),
(pumps, "set_pump", "get_pump", "_set_pump"),
(valves, "set_valve", "get_valve", "_set_valve"),
],
)
def test_element_patch_reads_after_shared_model_lock(
monkeypatch,
module,
setter_name: str,
getter_name: str,
builder_name: str,
) -> None:
events: list[str] = []
current = {"id": "E1"}
def fake_getter(_name: str, _id: str):
events.append("read")
return current
def fake_builder(_name: str, _changes: ChangeSet, supplied_current):
events.append("build")
assert supplied_current is current
return database.DatabaseCommand("UPDATE network.nodes SET id = id", [])
def fake_locked_command(_name: str, builder):
events.append("lock")
assert builder() is not None
return ChangeSet()
monkeypatch.setattr(module, getter_name, fake_getter)
monkeypatch.setattr(module, builder_name, fake_builder)
monkeypatch.setattr(module, "execute_locked_command", fake_locked_command)
getattr(module, setter_name)(
"project_a",
ChangeSet({"operation": "update", "type": "element", "id": "E1"}),
)
assert events == ["lock", "read", "build"]
tests/unit/test_wndb_endpoint_boundaries.py
+54
View File
@@ -0,0 +1,54 @@
import ast
from pathlib import Path
def test_wndb_editor_routes_are_synchronous() -> None:
"""Sync psycopg-backed routes must let FastAPI schedule them in its thread pool."""
endpoints = Path(__file__).resolve().parents[2] / "app" / "api" / "v1" / "endpoints"
route_roots = [endpoints / "network", endpoints / "components"]
async_routes: list[str] = []
for root in route_roots:
for path in root.glob("*.py"):
tree = ast.parse(path.read_text(encoding="utf-8"))
for node in tree.body:
if not isinstance(node, ast.AsyncFunctionDef):
continue
is_route = any(
isinstance(decorator, ast.Call)
and isinstance(decorator.func, ast.Attribute)
and isinstance(decorator.func.value, ast.Name)
and decorator.func.value.id == "router"
for decorator in node.decorator_list
)
if is_route:
async_routes.append(f"{path.name}:{node.lineno}:{node.name}")
assert async_routes == []
def test_simulation_routes_are_synchronous() -> None:
"""EPANET and synchronous WNDB work must run in FastAPI's thread pool."""
path = (
Path(__file__).resolve().parents[2]
/ "app"
/ "api"
/ "v1"
/ "endpoints"
/ "simulation.py"
)
tree = ast.parse(path.read_text(encoding="utf-8"))
async_routes = [
node.name
for node in tree.body
if isinstance(node, ast.AsyncFunctionDef)
and any(
isinstance(decorator, ast.Call)
and isinstance(decorator.func, ast.Attribute)
and isinstance(decorator.func.value, ast.Name)
and decorator.func.value.id == "router"
for decorator in node.decorator_list
)
]
assert async_routes == []
tests/unit/test_wndb_importer.py
+66
View File
@@ -0,0 +1,66 @@
from contextlib import contextmanager
from pathlib import Path
from app.native.wndb.core.database import ChangeSet
from app.native.wndb.inp import importer
def test_import_inp_uses_unique_file_and_always_removes_it(monkeypatch) -> None:
paths: list[str] = []
def fake_read_inp(_project: str, path: str, _version: str) -> bool:
assert Path(path).read_text(encoding="utf-8") == "[TITLE]\nmodel"
paths.append(path)
return True
monkeypatch.setattr(importer, "read_inp", fake_read_inp)
change_set = ChangeSet({"inp": "[TITLE]\nmodel"})
assert importer.import_inp("project_a", change_set) is True
assert importer.import_inp("project_a", change_set) is True
assert paths[0] != paths[1]
assert all(not Path(path).exists() for path in paths)
def test_read_inp_refreshes_after_committed_replace_when_cleanup_fails(
monkeypatch,
) -> None:
events: list[str] = []
@contextmanager
def fake_transaction(_project: str):
yield
monkeypatch.setattr(importer, "have_project", lambda _project: True)
monkeypatch.setattr(
importer, "temporary_project_name", lambda _project, _purpose: "staging"
)
monkeypatch.setattr(
importer,
"copy_project",
lambda _source, _target: events.append("copy"),
)
monkeypatch.setattr(importer, "project_transaction", fake_transaction)
monkeypatch.setattr(
importer, "parse_file", lambda *_args: events.append("parse")
)
monkeypatch.setattr(
importer,
"replace_project_model",
lambda *_args: events.append("replace"),
)
def fail_cleanup(_project: str) -> None:
events.append("cleanup-failed")
raise RuntimeError("drop failed")
monkeypatch.setattr(importer, "delete_project", fail_cleanup)
monkeypatch.setattr(
importer,
"refresh_materialized_views_after_commit",
lambda _project: events.append("refresh"),
)
assert importer.read_inp("project_a", "model.inp") is True
assert events == ["copy", "parse", "replace", "cleanup-failed", "refresh"]
tests/unit/test_wndb_materialized_refresh.py
+21 -2
View File
@@ -1,3 +1,5 @@
import pytest
from app.native.wndb.core import database
@@ -29,7 +31,7 @@ def test_direct_model_write_refreshes_materialized_views(monkeypatch) -> None:
)
monkeypatch.setattr(
database,
"refresh_materialized_views",
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
@@ -61,7 +63,7 @@ def test_batch_model_write_defers_materialized_view_refresh(monkeypatch) -> None
)
monkeypatch.setattr(
database,
"refresh_materialized_views",
"refresh_materialized_views_after_commit",
lambda _name: events.append("refresh"),
)
@@ -73,6 +75,23 @@ def test_batch_model_write_defers_materialized_view_refresh(monkeypatch) -> None
assert events == ["write"]
def test_refresh_after_commit_reports_that_changes_are_durable(monkeypatch) -> None:
def fail_refresh(_name: str) -> None:
raise RuntimeError("refresh failed")
monkeypatch.setattr(database, "refresh_materialized_views", fail_refresh)
with pytest.raises(
database.MaterializedViewRefreshAfterCommitError,
match="changes were committed",
) as exc_info:
database.refresh_materialized_views_after_commit("project_a")
assert exc_info.value.project == "project_a"
assert exc_info.value.changes_committed is True
assert isinstance(exc_info.value.__cause__, RuntimeError)
def test_non_gis_model_write_does_not_refresh_materialized_views(monkeypatch) -> None:
events: list[str] = []
monkeypatch.setattr(
tests/unit/test_wndb_network_views.py
+163
View File
@@ -0,0 +1,163 @@
from app.native.wndb.gis import network_views, region_geometry
from app.native.wndb.model import elements
def test_network_node_coords_use_one_unified_view_query(monkeypatch) -> None:
calls = []
def fake_read_all(name, statement, params=None):
calls.append((name, statement, params))
return [
{"id": "J-1", "x": 10.5, "y": 20.5, "node_type": "junction"},
{"id": "R-1", "x": 30.0, "y": 40.0, "node_type": "reservoir"},
]
monkeypatch.setattr(network_views, "read_all", fake_read_all)
assert network_views.get_network_node_coords("project_a") == {
"J-1": {"x": 10.5, "y": 20.5, "type": "junction"},
"R-1": {"x": 30.0, "y": 40.0, "type": "reservoir"},
}
assert len(calls) == 1
assert "FROM gis.network_nodes" in calls[0][1]
def test_network_link_nodes_use_one_unified_view_query(monkeypatch) -> None:
calls = []
def fake_read_all(name, statement, params=None):
calls.append((name, statement, params))
return [
{
"id": "P-1",
"link_type": "pipe",
"start_node_id": "J-1",
"end_node_id": "J-2",
},
{
"id": "PU-1",
"link_type": "pump",
"start_node_id": "R-1",
"end_node_id": "J-1",
},
]
monkeypatch.setattr(network_views, "read_all", fake_read_all)
assert network_views.get_network_link_nodes("project_a") == [
"P-1:pipe:J-1:J-2",
"PU-1:pump:R-1:J-1",
]
assert len(calls) == 1
assert "FROM gis.network_links" in calls[0][1]
def test_topology_rows_are_loaded_in_two_batch_queries(monkeypatch) -> None:
calls = []
responses = [
[{"id": "J-1", "x": 1.0, "y": 2.0, "node_type": "junction"}],
[
{
"id": "P-1",
"start_node_id": "J-1",
"end_node_id": "J-2",
"length": 12.5,
}
],
]
def fake_read_all(name, statement, params=None):
calls.append((name, statement, params))
return responses[len(calls) - 1]
monkeypatch.setattr(network_views, "read_all", fake_read_all)
nodes, links = network_views.get_topology_rows("project_a", ["J-1", "J-2"])
assert nodes == responses[0]
assert links == responses[1]
assert len(calls) == 2
assert "FROM network.nodes" in calls[0][1]
assert "FROM network.links" in calls[1][1]
assert "LEFT JOIN network.pipes" in calls[1][1]
def test_topology_builds_adjacency_from_batch_rows(monkeypatch) -> None:
monkeypatch.setattr(
region_geometry,
"get_topology_rows",
lambda _name, _node_ids: (
[
{"id": "J-1", "x": 1.0, "y": 2.0, "node_type": "junction"},
{"id": "R-1", "x": 3.0, "y": 4.0, "node_type": "reservoir"},
],
[
{
"id": "P-1",
"start_node_id": "J-1",
"end_node_id": "R-1",
"length": 20.0,
}
],
),
)
topology = region_geometry.Topology("project_a", ["J-1", "R-1"])
assert topology.max_x_node() == "R-1"
assert topology.nodes()["J-1"] == {
"x": 1.0,
"y": 2.0,
"type": "junction",
"links": ["P-1"],
}
assert topology.links()["P-1"] == {
"node1": "J-1",
"node2": "R-1",
"length": 20.0,
}
def test_junction_demands_are_mapped_from_authoritative_table(monkeypatch) -> None:
monkeypatch.setattr(
network_views,
"read_all",
lambda _name, _statement, _params: [
{
"junction_id": "J-1",
"sequence_no": 0,
"base_demand": 3.5,
"pattern_id": "PAT-1",
"category": None,
}
],
)
assert network_views.get_junction_demands("project_a", ["J-1"]) == {
"J-1": [
{"demand": 3.5, "pattern": "PAT-1", "category": None}
]
}
def test_all_node_links_scan_the_unified_view_once(monkeypatch) -> None:
calls = []
def fake_read_all_typed(name, statement, params):
calls.append((name, statement, params))
return [
{
"id": "P-1",
"start_node_id": "J-1",
"end_node_id": "J-2",
}
]
monkeypatch.setattr(elements, "read_all_typed", fake_read_all_typed)
assert elements.get_all_node_links("project_a") == {
"J-1": ["P-1"],
"J-2": ["P-1"],
}
assert len(calls) == 1
assert "FROM gis.network_links" in calls[0][1]
tests/unit/test_wndb_projects.py
+279
View File
@@ -0,0 +1,279 @@
from contextlib import contextmanager
import pytest
from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing
from app.native.wndb.core import projects
class _FakeCursor:
def __init__(self, *, rows=None, current_database: str = "postgres") -> None:
self.rows = rows or []
self.current_database = current_database
self.calls: list[tuple[object, object]] = []
self._last_statement = None
def __enter__(self):
return self
def __exit__(self, *_args):
return None
def __iter__(self):
return iter(self.rows)
def execute(self, statement, params=None):
self.calls.append((statement, params))
self._last_statement = statement
return self
def fetchone(self):
if isinstance(self._last_statement, str) and self._last_statement.startswith(
"select datallowconn"
):
return {"datallowconn": False}
return {"current_database": self.current_database}
class _FakeConnection:
def __init__(self, cursor: _FakeCursor) -> None:
self._cursor = cursor
def cursor(self, **_kwargs):
return self._cursor
def _admin_connection(cursor: _FakeCursor):
@contextmanager
def connection():
yield _FakeConnection(cursor)
return connection
@pytest.mark.parametrize(
"name",
[
"postgres",
"project",
"system_hub",
"SYSTEM_HUB",
"tjwater_v2_template",
"another_template",
],
)
def test_delete_project_rejects_protected_database_before_side_effects(
monkeypatch, name
) -> None:
closed: list[str] = []
monkeypatch.setattr(projects, "close_project_pool", closed.append)
monkeypatch.setattr(
projects,
"admin_connection",
lambda: pytest.fail("protected database opened an administration connection"),
)
with pytest.raises(ValueError, match="protected"):
projects.delete_project(name)
assert closed == []
def test_copy_project_rejects_metadata_source(monkeypatch) -> None:
monkeypatch.setattr(
projects,
"admin_connection",
lambda: pytest.fail("protected database opened an administration connection"),
)
with pytest.raises(ValueError, match="protected"):
projects.copy_project("system_hub", "copy")
def test_copy_project_rejects_unconfigured_template_source(monkeypatch) -> None:
monkeypatch.setattr(
projects,
"admin_connection",
lambda: pytest.fail("unconfigured template opened an administration connection"),
)
with pytest.raises(ValueError, match="protected"):
projects.copy_project("tjwater_next_template", "copy")
def test_create_project_allows_the_protected_template_as_source(monkeypatch) -> None:
cursor = _FakeCursor()
closed: list[str] = []
monkeypatch.setattr(projects, "admin_connection", _admin_connection(cursor))
monkeypatch.setattr(projects, "close_project_pool", closed.append)
projects.create_project("project_a")
assert closed == ["tjwater_v2_template", "project_a"]
assert any(call[1] == ("tjwater_v2_template",) for call in cursor.calls)
assert any(
isinstance(call[0], str)
and call[0].startswith("select pg_terminate_backend")
and call[1] == ("tjwater_v2_template",)
for call in cursor.calls
)
assert any("create database" in str(call[0]).lower() for call in cursor.calls)
def test_list_project_excludes_metadata_database(monkeypatch) -> None:
cursor = _FakeCursor(rows=[{"datname": "project_a"}])
monkeypatch.setattr(projects, "admin_connection", _admin_connection(cursor))
assert projects.list_project() == ["project_a"]
excluded = cursor.calls[0][1][0]
assert "system_hub" in excluded
assert "project" in excluded
assert "tjwater_v2_template" in excluded
def test_delete_project_uses_routed_physical_database_name(monkeypatch) -> None:
cursor = _FakeCursor()
closed: list[str] = []
monkeypatch.setattr(projects, "admin_connection", _admin_connection(cursor))
monkeypatch.setattr(projects, "close_project_pool", closed.append)
routing = ActiveProjectRouting(
project_code="logical-project",
business_dsn="postgresql://user:password@db.example/tjwater_v2",
)
with activate_project_routing(routing):
projects.delete_project("logical-project")
assert closed == ["logical-project"]
assert any(call[1] == ("tjwater_v2",) for call in cursor.calls)
def test_temporary_project_names_are_unique_and_postgres_safe() -> None:
first = projects.temporary_project_name(
"TJWater V2 / Production With A Very Long Project Name",
"Burst Analysis",
)
second = projects.temporary_project_name(
"TJWater V2 / Production With A Very Long Project Name",
"Burst Analysis",
)
assert first != second
assert len(first.encode("utf-8")) <= 63
assert first.startswith("tjw_tmp_burst_analysis_tjwat")
def test_temporary_database_capacity_rejects_creation_at_limit(
monkeypatch,
) -> None:
cursor = _FakeCursor()
monkeypatch.setattr(projects.settings, "WNDB_TEMP_DB_MAX_COUNT", 2)
cursor.fetchone = lambda: {"count": 2}
with pytest.raises(RuntimeError, match="limit reached"):
with projects._temporary_database_capacity(cursor, "tjw_tmp_analysis_123"):
pytest.fail("capacity guard yielded after reaching the limit")
assert any("pg_advisory_unlock" in str(statement) for statement, _ in cursor.calls)
def test_temporary_project_database_cleans_up_after_failure(monkeypatch) -> None:
calls: list[tuple[str, ...]] = []
monkeypatch.setattr(
projects,
"temporary_project_name",
lambda project, purpose: "isolated_run",
)
monkeypatch.setattr(
projects,
"copy_project",
lambda source, target: calls.append(("copy", source, target)),
)
monkeypatch.setattr(projects, "have_project", lambda name: True)
monkeypatch.setattr(
projects,
"delete_project",
lambda name: calls.append(("delete", name)),
)
monkeypatch.setattr(
"app.native.wndb.core.model_replace.replace_project_model",
lambda target, source, *, copy_source_scada: calls.append(
("clone", target, source, str(copy_source_scada))
),
)
monkeypatch.setattr(
"app.native.wndb.core.database.refresh_materialized_views_after_commit",
lambda name: calls.append(("refresh", name)),
)
with pytest.raises(RuntimeError, match="analysis failed"):
with projects.temporary_project_database("project_a", "age") as name:
assert name == "isolated_run"
raise RuntimeError("analysis failed")
assert calls == [
("copy", "tjwater_v2_template", "isolated_run"),
("clone", "isolated_run", "project_a", "True"),
("refresh", "isolated_run"),
("delete", "isolated_run"),
]
def test_temporary_template_database_does_not_clone_a_project(monkeypatch) -> None:
calls: list[tuple[str, ...]] = []
monkeypatch.setattr(
projects,
"temporary_project_name",
lambda project, purpose: "empty_conversion",
)
monkeypatch.setattr(
projects,
"copy_project",
lambda source, target: calls.append(("copy", source, target)),
)
monkeypatch.setattr(projects, "have_project", lambda name: True)
monkeypatch.setattr(
projects,
"delete_project",
lambda name: calls.append(("delete", name)),
)
with projects.temporary_template_database("conversion", "v3_to_v2") as name:
assert name == "empty_conversion"
assert calls == [
("copy", "tjwater_v2_template", "empty_conversion"),
("delete", "empty_conversion"),
]
def test_clean_project_deletes_only_explicit_unique_targets(monkeypatch) -> None:
cursor = _FakeCursor()
closed: list[str] = []
monkeypatch.setattr(projects, "admin_connection", _admin_connection(cursor))
monkeypatch.setattr(projects, "close_project_pool", closed.append)
projects.clean_project(["temp_a", "temp_b", "temp_a"])
assert closed == ["temp_a", "temp_b"]
termination_targets = [
params[0]
for statement, params in cursor.calls
if isinstance(statement, str) and statement.startswith("select pg_terminate_backend")
]
assert termination_targets == ["temp_a", "temp_b"]
def test_clean_project_validates_all_targets_before_deleting(monkeypatch) -> None:
closed: list[str] = []
monkeypatch.setattr(projects, "close_project_pool", closed.append)
monkeypatch.setattr(
projects,
"admin_connection",
lambda: pytest.fail("invalid targets opened an administration connection"),
)
with pytest.raises(ValueError, match="protected"):
projects.clean_project(["temp_a", "system_hub"])
assert closed == []
tests/unit/test_wndb_query_safety.py
+21
View File
@@ -3,6 +3,7 @@ from pathlib import Path
from app.infra.db.postgresql import scada_assets
from app.native.wndb.core.database import ChangeSet, sql_literal
from app.native.wndb.gis import coordinates
from app.native.wndb.model import controls, junctions, patterns
@@ -53,6 +54,26 @@ def test_get_all_scada_info_reads_materialized_view(monkeypatch) -> None:
assert "FROM gis.scada_devices" in statements[0]
def test_get_missing_node_coord_does_not_write_default_geometry(monkeypatch) -> None:
calls: list[tuple[str, str, tuple[str, ...]]] = []
def fake_try_read(name, statement, params):
calls.append((name, statement, params))
return None
monkeypatch.setattr(coordinates, "try_read", fake_try_read)
assert coordinates.get_node_coord("project_a", "J-1") == {"x": 0.0, "y": 0.0}
assert calls == [
(
"project_a",
"select st_astext(geom) as coord_geom "
"from gis.node_geometries where node_id = %s",
("J-1",),
)
]
def test_sql_literal_keeps_attacker_text_inside_one_postgres_literal() -> None:
malicious = "x'); DROP SCHEMA network CASCADE; --"