refactor(db)!: adopt project-routed pooled databases
Reorganize WNDB by responsibility and remove legacy scheme endpoints.\n\nRoute analysis and time-series access through project pools, preserve transactional realtime replacement, and refresh GIS materialized views after writes.\n\nAdd database architecture documentation, live pooling coverage, API contract updates, and executable container verification.\n\nBREAKING CHANGE: legacy scheme APIs and flat app.native.wndb module imports are removed.
This commit is contained in:
@@ -10,12 +10,10 @@ from fastapi import APIRouter, FastAPI, Query
|
||||
from fastapi.routing import APIRoute
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1.endpoints import schemes as schemes_endpoint
|
||||
from app.api.v1.endpoints import simulation as simulation_endpoint
|
||||
from app.api.pagination import PaginatedList
|
||||
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.metadata_dependencies import get_current_metadata_user
|
||||
from app.auth.project_dependencies import (
|
||||
ProjectContext,
|
||||
get_project_business_routing,
|
||||
@@ -173,12 +171,6 @@ def test_rest_contract_uses_header_project_context() -> None:
|
||||
assert "network" not in schema.get("properties", {})
|
||||
assert "network_name" not in schema.get("properties", {})
|
||||
|
||||
placement_schema = document["components"]["schemas"][
|
||||
"PressureSensorPlacementRest"
|
||||
]
|
||||
assert "name" not in placement_schema["properties"]
|
||||
assert "username" not in placement_schema["properties"]
|
||||
|
||||
assert "/api/v1/burst-analysis" not in document["paths"]
|
||||
assert "/api/v1/getpipeproperties/" not in document["paths"]
|
||||
|
||||
@@ -294,54 +286,10 @@ def test_sensor_placement_excel_export_is_post() -> None:
|
||||
if isinstance(route, APIRoute)
|
||||
}
|
||||
assert methods_by_path[
|
||||
"/sensor-placement-schemes/{scheme_id}/exports/excel"
|
||||
"/sensor-placement-runs/{run_id}/exports/excel"
|
||||
] == {"POST"}
|
||||
|
||||
|
||||
def test_rest_runtime_consumes_injected_project_context(monkeypatch) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_get_all_schemes(network, scheme_type=None, query_date=None):
|
||||
captured.update(
|
||||
network=network,
|
||||
scheme_type=scheme_type,
|
||||
query_date=query_date,
|
||||
business_dsn=get_project_pgconn_string(network),
|
||||
)
|
||||
return [{"scheme_name": "burst_case", "scheme_type": scheme_type}]
|
||||
|
||||
monkeypatch.setattr(
|
||||
schemes_endpoint,
|
||||
"get_all_schemes",
|
||||
fake_get_all_schemes,
|
||||
)
|
||||
app = FastAPI(redirect_slashes=False)
|
||||
app.include_router(api_router, prefix="/api/v1")
|
||||
project_context = ProjectContext(
|
||||
project_id=uuid4(),
|
||||
project_code="fengyang",
|
||||
user_id=uuid4(),
|
||||
project_role="viewer",
|
||||
)
|
||||
_override_project_routing(app, project_context)
|
||||
|
||||
response = TestClient(app, raise_server_exceptions=False).get(
|
||||
"/api/v1/schemes",
|
||||
params={"scheme_type": "burst_analysis"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"network": "fengyang",
|
||||
"scheme_type": "burst_analysis",
|
||||
"query_date": None,
|
||||
"business_dsn": "postgresql://user:password@biz/fengyang",
|
||||
}
|
||||
assert response.json()["items"] == [
|
||||
{"scheme_name": "burst_case", "scheme_type": "burst_analysis"}
|
||||
]
|
||||
|
||||
|
||||
def test_rest_runtime_wraps_handler_paginated_list() -> None:
|
||||
source_router = APIRouter()
|
||||
|
||||
@@ -370,51 +318,6 @@ def test_rest_runtime_wraps_handler_paginated_list() -> None:
|
||||
}
|
||||
|
||||
|
||||
def test_sensor_placement_body_uses_authenticated_project_and_user(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_pressure_sensor_placement_kmeans(**kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(
|
||||
simulation_endpoint,
|
||||
"pressure_sensor_placement_kmeans",
|
||||
fake_pressure_sensor_placement_kmeans,
|
||||
)
|
||||
app = FastAPI(redirect_slashes=False)
|
||||
app.include_router(api_router, prefix="/api/v1")
|
||||
project_context = ProjectContext(
|
||||
project_id=uuid4(),
|
||||
project_code="project_a",
|
||||
user_id=uuid4(),
|
||||
project_role="member",
|
||||
)
|
||||
_override_project_routing(app, project_context)
|
||||
app.dependency_overrides[get_current_metadata_user] = lambda: type(
|
||||
"User", (), {"username": "alice"}
|
||||
)()
|
||||
|
||||
response = TestClient(app, raise_server_exceptions=False).post(
|
||||
"/api/v1/pressure-sensor-placement-kmeans",
|
||||
json={
|
||||
"scheme_name": "placement_01",
|
||||
"sensor_number": 5,
|
||||
"min_diameter": 100,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"name": "project_a",
|
||||
"scheme_name": "placement_01",
|
||||
"sensor_number": 5,
|
||||
"min_diameter": 100,
|
||||
"username": "alice",
|
||||
}
|
||||
|
||||
|
||||
def test_rest_runtime_json_encodes_untyped_datetime_response() -> None:
|
||||
source_router = APIRouter()
|
||||
|
||||
|
||||
@@ -42,11 +42,10 @@ def _load_project_module(monkeypatch):
|
||||
"read_inp": lambda network, inp: True,
|
||||
"dump_inp": lambda network, inp: True,
|
||||
"get_all_vertices": lambda network: [],
|
||||
"get_all_scada_elements": lambda network: [],
|
||||
"get_all_scada_info": lambda network: [],
|
||||
"get_all_district_metering_areas": lambda network: [],
|
||||
"get_all_service_areas": lambda network: [],
|
||||
"get_all_virtual_districts": lambda network: [],
|
||||
"get_extension_data": lambda network, key: None,
|
||||
"convert_inp_v3_to_v2": lambda inp: DummyChangeSet({"inp": inp}),
|
||||
},
|
||||
)
|
||||
@@ -55,16 +54,6 @@ def _load_project_module(monkeypatch):
|
||||
"app.auth.project_dependencies",
|
||||
{"get_metadata_repository": lambda: None},
|
||||
)
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.infra.db.postgresql.database",
|
||||
{"get_database_instance": lambda network: None},
|
||||
)
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.infra.db.timescaledb.database",
|
||||
{"get_database_instance": lambda network: None},
|
||||
)
|
||||
return load_module_from_path(
|
||||
"tests_project_endpoints_module",
|
||||
"app/api/v1/endpoints/project.py",
|
||||
@@ -109,16 +98,12 @@ def test_project_info_returns_project_workspace(monkeypatch):
|
||||
assert "geoserver" not in payload
|
||||
|
||||
|
||||
def test_open_project_returns_network_even_when_db_connection_fails(monkeypatch):
|
||||
def test_open_project_uses_unified_wndb_connection_path(monkeypatch):
|
||||
module = _load_project_module(monkeypatch)
|
||||
called = []
|
||||
|
||||
monkeypatch.setattr(module, "open_project", lambda network: called.append(network))
|
||||
|
||||
async def failing_get_pg_db(network):
|
||||
raise RuntimeError("db down")
|
||||
|
||||
monkeypatch.setattr(module, "get_pg_db", failing_get_pg_db)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
|
||||
response = client.post("/api/v1/projects/current", params={"network": "demo"})
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from typing import Any
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from tests.conftest import build_test_app, install_stub, load_module_from_path
|
||||
@@ -7,16 +5,7 @@ from tests.conftest import build_test_app, install_stub, load_module_from_path
|
||||
|
||||
class DummyChangeSet:
|
||||
def __init__(self, operations=None):
|
||||
if operations is None:
|
||||
self.operations = []
|
||||
elif isinstance(operations, dict):
|
||||
self.operations = [operations]
|
||||
else:
|
||||
self.operations = operations
|
||||
|
||||
|
||||
def _noop(*args, **kwargs):
|
||||
return None
|
||||
self.operations = [operations] if isinstance(operations, dict) else operations or []
|
||||
|
||||
|
||||
def _load_regions_module(monkeypatch):
|
||||
@@ -25,41 +14,18 @@ def _load_regions_module(monkeypatch):
|
||||
monkeypatch,
|
||||
"app.services.tjnetwork",
|
||||
{
|
||||
"Any": Any,
|
||||
"ChangeSet": DummyChangeSet,
|
||||
"add_district_metering_area": _noop,
|
||||
"add_region": _noop,
|
||||
"add_service_area": _noop,
|
||||
"add_virtual_district": _noop,
|
||||
"calculate_district_metering_area_for_network": lambda *args, **kwargs: [],
|
||||
"calculate_district_metering_area_for_nodes": lambda *args, **kwargs: [],
|
||||
"calculate_district_metering_area_for_region": lambda *args, **kwargs: [],
|
||||
"calculate_service_area": lambda network: [],
|
||||
"calculate_virtual_district": lambda *args, **kwargs: {},
|
||||
"delete_district_metering_area": _noop,
|
||||
"delete_region": _noop,
|
||||
"delete_service_area": _noop,
|
||||
"delete_virtual_district": _noop,
|
||||
"generate_district_metering_area": _noop,
|
||||
"generate_service_area": _noop,
|
||||
"generate_sub_district_metering_area": _noop,
|
||||
"generate_virtual_district": _noop,
|
||||
"get_all_district_metering_area_ids": lambda network: [],
|
||||
"get_all_district_metering_areas": lambda network: [],
|
||||
"get_all_service_areas": lambda network: [],
|
||||
"get_all_virtual_districts": lambda network: [],
|
||||
"get_district_metering_area": lambda network, area_id: {},
|
||||
"get_district_metering_area_schema": lambda network: {},
|
||||
"get_region": lambda network, region_id: {},
|
||||
"get_region_schema": lambda network: {},
|
||||
"get_service_area": lambda network, area_id: {},
|
||||
"get_service_area_schema": lambda network: {},
|
||||
"get_virtual_district": lambda network, area_id: {},
|
||||
"get_virtual_district_schema": lambda network: {},
|
||||
"set_district_metering_area": _noop,
|
||||
"set_region": _noop,
|
||||
"set_service_area": _noop,
|
||||
"set_virtual_district": _noop,
|
||||
"add_region": lambda network, cs: cs,
|
||||
"delete_region": lambda network, cs: cs,
|
||||
"get_nodes_in_region": lambda network, region_id: ["J1"],
|
||||
"get_region": lambda network, region_id: {
|
||||
"id": region_id,
|
||||
"region_type": "DMA",
|
||||
"boundary": [[0, 0], [1, 0], [0, 0]],
|
||||
},
|
||||
"get_region_schema": lambda network: {"id": {"type": "str"}},
|
||||
"get_regions": lambda network: ["DMA-1"],
|
||||
"set_region": lambda network, cs: cs,
|
||||
},
|
||||
)
|
||||
return load_module_from_path(
|
||||
@@ -68,87 +34,47 @@ def _load_regions_module(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_removed_routes_are_absent_and_return_404(monkeypatch):
|
||||
def test_regions_are_exposed_as_one_generic_resource(monkeypatch):
|
||||
module = _load_regions_module(monkeypatch)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
|
||||
openapi = client.get("/openapi.json").json()
|
||||
|
||||
assert "/api/v1/calculateregion/" not in openapi["paths"]
|
||||
assert "/api/v1/getallregions/" not in openapi["paths"]
|
||||
assert "/api/v1/generateregion/" not in openapi["paths"]
|
||||
assert "/api/v1/calculatedistrictmeteringarea/" not in openapi["paths"]
|
||||
assert client.get("/api/v1/calculateregion/", params={"network": "demo", "time_index": 0}).status_code == 404
|
||||
assert client.get("/api/v1/calculatedistrictmeteringarea/", params={"network": "demo"}).status_code == 404
|
||||
|
||||
|
||||
def test_calculate_service_area_contract_uses_only_network(monkeypatch):
|
||||
module = _load_regions_module(monkeypatch)
|
||||
calls = []
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"calculate_service_area",
|
||||
lambda network: calls.append(network) or [{"source-1": ["n1", "n2"]}],
|
||||
)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
|
||||
response = client.post(
|
||||
"/api/v1/service-area-calculations",
|
||||
params={"network": "demo", "time_index": 5},
|
||||
)
|
||||
schema = client.get("/openapi.json").json()
|
||||
response = client.get("/api/v1/regions", params={"network": "demo"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == [{"source-1": ["n1", "n2"]}]
|
||||
assert calls == ["demo"]
|
||||
parameter_names = [
|
||||
item["name"]
|
||||
for item in schema["paths"]["/api/v1/service-area-calculations"]["post"]["parameters"]
|
||||
]
|
||||
assert parameter_names == ["network"]
|
||||
assert response.json()[0]["region_type"] == "DMA"
|
||||
|
||||
|
||||
def test_add_district_metering_area_converts_boundary_to_tuples(monkeypatch):
|
||||
def test_add_region_converts_boundary_to_tuples(monkeypatch):
|
||||
module = _load_regions_module(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
def fake_add(network, change_set):
|
||||
captured["network"] = network
|
||||
captured["boundary"] = change_set.operations[0]["boundary"]
|
||||
return {"ok": True}
|
||||
def add(network, changeset):
|
||||
captured["operation"] = changeset.operations[0]
|
||||
return changeset
|
||||
|
||||
monkeypatch.setattr(module, "add_district_metering_area", fake_add)
|
||||
monkeypatch.setattr(module, "add_region", add)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
|
||||
response = client.post(
|
||||
"/api/v1/district-metering-areas",
|
||||
"/api/v1/regions",
|
||||
params={"network": "demo"},
|
||||
json={"id": "dma-1", "boundary": [[1, 2], [3, 4], [1, 2]]},
|
||||
json={
|
||||
"id": "DMA-1",
|
||||
"region_type": "DMA",
|
||||
"boundary": [[0, 0], [1, 0], [0, 0]],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"network": "demo",
|
||||
"boundary": [(1, 2), (3, 4), (1, 2)],
|
||||
}
|
||||
assert captured["operation"]["boundary"] == [(0, 0), (1, 0), (0, 0)]
|
||||
|
||||
|
||||
def test_generate_virtual_district_reads_centers_from_body(monkeypatch):
|
||||
def test_region_nodes_use_generic_region_id(monkeypatch):
|
||||
module = _load_regions_module(monkeypatch)
|
||||
captured = {}
|
||||
|
||||
def fake_generate(network, centers, inflate_delta):
|
||||
captured["args"] = (network, centers, inflate_delta)
|
||||
return {"generated": True}
|
||||
|
||||
monkeypatch.setattr(module, "generate_virtual_district", fake_generate)
|
||||
client = TestClient(build_test_app(module.router, "/api/v1"))
|
||||
|
||||
response = client.post(
|
||||
"/api/v1/virtual-district-generation-runs",
|
||||
params={"network": "demo", "inflate_delta": 0.75},
|
||||
json={"centers": ["J1", "J2"]},
|
||||
response = client.get(
|
||||
"/api/v1/regions/nodes", params={"network": "demo", "id": "DMA-1"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["args"] == ("demo", ["J1", "J2"], 0.75)
|
||||
assert response.json() == ["J1"]
|
||||
|
||||
@@ -1,102 +0,0 @@
|
||||
from datetime import date
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1.endpoints import schemes as schemes_endpoint
|
||||
|
||||
|
||||
def _build_client() -> TestClient:
|
||||
app = FastAPI()
|
||||
app.include_router(schemes_endpoint.router, prefix="/api/v1")
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_get_schemes_forwards_optional_scheme_type(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_get_all_schemes(network, scheme_type=None, query_date=None):
|
||||
captured["network"] = network
|
||||
captured["scheme_type"] = scheme_type
|
||||
captured["query_date"] = query_date
|
||||
return [
|
||||
{
|
||||
"scheme_id": 1,
|
||||
"scheme_name": "burst_case",
|
||||
"scheme_type": scheme_type,
|
||||
}
|
||||
]
|
||||
|
||||
monkeypatch.setattr(schemes_endpoint, "get_all_schemes", fake_get_all_schemes)
|
||||
|
||||
response = _build_client().get(
|
||||
"/api/v1/schemes",
|
||||
params={"network": "demo", "scheme_type": "burst_analysis"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"network": "demo",
|
||||
"scheme_type": "burst_analysis",
|
||||
"query_date": None,
|
||||
}
|
||||
assert response.json()[0]["scheme_type"] == "burst_analysis"
|
||||
|
||||
|
||||
def test_get_schemes_forwards_query_date(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_get_all_schemes(network, scheme_type=None, query_date=None):
|
||||
captured["network"] = network
|
||||
captured["scheme_type"] = scheme_type
|
||||
captured["query_date"] = query_date
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(schemes_endpoint, "get_all_schemes", fake_get_all_schemes)
|
||||
|
||||
response = _build_client().get(
|
||||
"/api/v1/schemes",
|
||||
params={
|
||||
"network": "demo",
|
||||
"scheme_type": "dma_leak_identification",
|
||||
"query_date": "2026-01-02T00:00:00+08:00",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"network": "demo",
|
||||
"scheme_type": "dma_leak_identification",
|
||||
"query_date": date(2026, 1, 2),
|
||||
}
|
||||
|
||||
|
||||
def test_get_scheme_detail_forwards_scheme_type(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_query_scheme_detail(name, scheme_name, scheme_type=None):
|
||||
captured["name"] = name
|
||||
captured["scheme_name"] = scheme_name
|
||||
captured["scheme_type"] = scheme_type
|
||||
return {
|
||||
"scheme_name": scheme_name,
|
||||
"scheme_type": scheme_type,
|
||||
"rows": [{"Area": "1", "LeakageFlow_m3_per_s": 0.1}],
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
schemes_endpoint, "query_scheme_detail", fake_query_scheme_detail
|
||||
)
|
||||
|
||||
response = _build_client().get(
|
||||
"/api/v1/schemes/dma_001",
|
||||
params={"network": "demo", "scheme_type": "dma_leak_identification"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured == {
|
||||
"name": "demo",
|
||||
"scheme_name": "dma_001",
|
||||
"scheme_type": "dma_leak_identification",
|
||||
}
|
||||
assert response.json()["scheme_name"] == "dma_001"
|
||||
@@ -1,424 +1,124 @@
|
||||
from datetime import datetime, timezone
|
||||
from io import BytesIO
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from tests.conftest import build_test_app, install_stub, load_module_from_path
|
||||
from app.api.v1.endpoints import sensor_placement as endpoint
|
||||
from tests.conftest import build_test_app
|
||||
|
||||
|
||||
class NotFoundError(LookupError):
|
||||
pass
|
||||
RUN_ID = uuid4()
|
||||
|
||||
|
||||
class ValidationError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class ConflictError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _scheme(**overrides):
|
||||
def _run(**overrides):
|
||||
value = {
|
||||
"id": 7,
|
||||
"scheme_name": "北区测压点",
|
||||
"sensor_number": 2,
|
||||
"run_id": RUN_ID,
|
||||
"name": "北区测压点",
|
||||
"sensor_count": 1,
|
||||
"min_diameter": 300,
|
||||
"username": "alice",
|
||||
"create_time": datetime(2026, 7, 30, 8, 0, tzinfo=timezone.utc),
|
||||
"sensor_location": ["J1", "J2"],
|
||||
"created_by": "alice",
|
||||
"created_at": datetime(2026, 8, 24, tzinfo=timezone.utc),
|
||||
"status": "completed",
|
||||
"sensor_locations": ["J1"],
|
||||
"sensor_points": [
|
||||
{
|
||||
"node_id": "J1",
|
||||
"max_pipe_diameter": 400.0,
|
||||
"project_x": 13500000.0,
|
||||
"project_y": 3600000.0,
|
||||
"map_x": 13500000.0,
|
||||
"map_y": 3600000.0,
|
||||
"project_x": 1.0,
|
||||
"project_y": 2.0,
|
||||
"map_x": 3.0,
|
||||
"map_y": 4.0,
|
||||
"longitude": 121.0,
|
||||
"latitude": 31.0,
|
||||
"elevation": 4.5,
|
||||
},
|
||||
{
|
||||
"node_id": "J2",
|
||||
"max_pipe_diameter": 300.0,
|
||||
"project_x": 13500100.0,
|
||||
"project_y": 3600100.0,
|
||||
"map_x": 13500100.0,
|
||||
"map_y": 3600100.0,
|
||||
"longitude": 121.001,
|
||||
"latitude": 31.001,
|
||||
"elevation": 5.0,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
value.update(overrides)
|
||||
return value
|
||||
|
||||
|
||||
def _load_module(monkeypatch):
|
||||
install_stub(monkeypatch, "app.algorithms", package=True)
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.algorithms.sensor",
|
||||
{
|
||||
"pressure_sensor_placement_kmeans": lambda **kwargs: {"id": 7},
|
||||
"pressure_sensor_placement_sensitivity": lambda **kwargs: {"id": 7},
|
||||
},
|
||||
def _client(project_role="member", username="alice", role="user"):
|
||||
app = build_test_app(endpoint.router, "/api/v1")
|
||||
app.dependency_overrides[endpoint.get_project_context] = lambda: SimpleNamespace(
|
||||
project_code="tjwater", project_role=project_role
|
||||
)
|
||||
install_stub(monkeypatch, "app.auth", package=True)
|
||||
|
||||
async def current_user():
|
||||
return SimpleNamespace(
|
||||
username="alice",
|
||||
role="user",
|
||||
is_superuser=False,
|
||||
)
|
||||
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.auth.metadata_dependencies",
|
||||
{"get_current_metadata_user": current_user},
|
||||
)
|
||||
|
||||
class ProjectContext:
|
||||
def __init__(self, project_code: str, project_role: str = "member"):
|
||||
self.project_code = project_code
|
||||
self.project_role = project_role
|
||||
|
||||
async def project_context():
|
||||
return ProjectContext("tjwater")
|
||||
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.auth.project_dependencies",
|
||||
{
|
||||
"ProjectContext": ProjectContext,
|
||||
"get_project_context": project_context,
|
||||
},
|
||||
)
|
||||
install_stub(monkeypatch, "app.services", package=True)
|
||||
install_stub(
|
||||
monkeypatch,
|
||||
"app.services.sensor_placement",
|
||||
{
|
||||
"SensorPlacementConflictError": ConflictError,
|
||||
"SensorPlacementNotFoundError": NotFoundError,
|
||||
"SensorPlacementValidationError": ValidationError,
|
||||
"build_sensor_placement_workbook": lambda **kwargs: BytesIO(b"xlsx"),
|
||||
"can_edit_sensor_placement": (
|
||||
lambda user, scheme: user.username == scheme["username"]
|
||||
or user.role == "admin"
|
||||
or user.is_superuser
|
||||
),
|
||||
"get_sensor_placement_scheme": lambda network, scheme_id: _scheme(
|
||||
id=scheme_id
|
||||
),
|
||||
"get_sensor_placement_candidate": (
|
||||
lambda network, node_id: _scheme()["sensor_points"][0]
|
||||
),
|
||||
"update_sensor_placement_scheme": (
|
||||
lambda network, scheme_id, **kwargs: _scheme(
|
||||
id=scheme_id,
|
||||
sensor_location=kwargs["sensor_location"],
|
||||
sensor_number=len(kwargs["sensor_location"]),
|
||||
)
|
||||
),
|
||||
},
|
||||
)
|
||||
return load_module_from_path(
|
||||
"tests_sensor_placement_endpoints_module",
|
||||
"app/api/v1/endpoints/sensor_placement.py",
|
||||
)
|
||||
|
||||
|
||||
def _client(module, user=None, project_role="member"):
|
||||
app = build_test_app(module.router, "/api/v1")
|
||||
if user is None:
|
||||
user = SimpleNamespace(
|
||||
username="alice",
|
||||
role="user",
|
||||
is_superuser=False,
|
||||
)
|
||||
app.dependency_overrides[module.get_current_metadata_user] = lambda: user
|
||||
app.dependency_overrides[module.get_project_context] = lambda: (
|
||||
module.ProjectContext("tjwater", project_role)
|
||||
app.dependency_overrides[endpoint.get_current_metadata_user] = lambda: SimpleNamespace(
|
||||
username=username, role=role, is_superuser=False
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_optimize_returns_created_scheme(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
def test_optimize_returns_analysis_run(monkeypatch):
|
||||
captured = {}
|
||||
monkeypatch.setattr(
|
||||
endpoint,
|
||||
"pressure_sensor_placement_kmeans",
|
||||
lambda **kwargs: captured.update(kwargs) or {"run_id": RUN_ID},
|
||||
)
|
||||
monkeypatch.setattr(endpoint, "get_sensor_placement_run", lambda *_: _run())
|
||||
|
||||
def optimize(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return {"id": 7}
|
||||
|
||||
monkeypatch.setattr(module, "pressure_sensor_placement_kmeans", optimize)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
response = _client().post(
|
||||
"/api/v1/sensor-placement-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": "北区测压点",
|
||||
"run_name": "北区测压点",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"sensor_count": 1,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["sensor_location"] == ["J1", "J2"]
|
||||
assert response.json()["run_id"] == str(RUN_ID)
|
||||
assert captured["username"] == "alice"
|
||||
|
||||
|
||||
def test_get_candidate_returns_maximum_incident_pipe_diameter(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
|
||||
response = _client(module).get(
|
||||
"/api/v1/sensor-placement-candidates/J1",
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["node_id"] == "J1"
|
||||
assert response.json()["max_pipe_diameter"] == 400.0
|
||||
|
||||
|
||||
def test_optimize_rejects_unsupported_sensor_type(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
def test_optimize_rejects_project_mismatch(monkeypatch):
|
||||
response = _client().post(
|
||||
"/api/v1/sensor-placement-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": "北区测流点",
|
||||
"sensor_type": "flow",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_optimize_rejects_network_outside_project_context(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
json={
|
||||
"network": "other_project",
|
||||
"scheme_name": "越权方案",
|
||||
"network": "other",
|
||||
"run_name": "越权运行",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
"sensor_count": 1,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_optimize_rejects_network_path_traversal(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
def test_viewer_cannot_update_run(monkeypatch):
|
||||
monkeypatch.setattr(endpoint, "get_sensor_placement_run", lambda *_: _run())
|
||||
|
||||
response = _client(project_role="viewer").put(
|
||||
f"/api/v1/sensor-placement-runs/{RUN_ID}",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"network": "../other_project",
|
||||
"scheme_name": "非法路径",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_optimize_rejects_unbounded_sensor_count(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": "超大方案",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 201,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_optimize_rejects_viewer_project_role(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module, project_role="viewer").post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": "只读成员方案",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
"expected_sensor_locations": ["J1"],
|
||||
"sensor_locations": ["J2"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"project_role",
|
||||
["owner", "admin", "modeler", "dispatcher", "auditor"],
|
||||
)
|
||||
def test_legacy_project_roles_cannot_optimize(monkeypatch, project_role):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module, project_role=project_role).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": f"{project_role}方案",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
def test_export_returns_xlsx(monkeypatch):
|
||||
monkeypatch.setattr(endpoint, "get_sensor_placement_run", lambda *_: _run())
|
||||
monkeypatch.setattr(
|
||||
endpoint,
|
||||
"build_sensor_placement_workbook",
|
||||
lambda **kwargs: BytesIO(b"xlsx"),
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_optimize_maps_running_project_job_to_409(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
|
||||
def conflict(**kwargs):
|
||||
raise ConflictError("当前项目已有监测点优化任务正在运行，请稍后重试")
|
||||
|
||||
monkeypatch.setattr(module, "pressure_sensor_placement_kmeans", conflict)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-optimization-runs",
|
||||
json={
|
||||
"network": "tjwater",
|
||||
"scheme_name": "并发方案",
|
||||
"sensor_type": "pressure",
|
||||
"method": "kmeans",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
|
||||
|
||||
def test_viewer_reads_scheme_as_non_editable(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module, project_role="viewer").get(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
response = _client().post(
|
||||
f"/api/v1/sensor-placement-runs/{RUN_ID}/exports/excel",
|
||||
params={"network": "tjwater"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["can_edit"] is False
|
||||
|
||||
|
||||
def test_update_rejects_non_owner(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(
|
||||
module,
|
||||
SimpleNamespace(username="bob", role="user", is_superuser=False),
|
||||
).put(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"expected_sensor_location": ["J1", "J2"],
|
||||
"sensor_location": ["J1", "J3"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_update_rejects_owner_with_viewer_project_role(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module, project_role="viewer").put(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"expected_sensor_location": ["J1", "J2"],
|
||||
"sensor_location": ["J1", "J3"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_admin_can_overwrite_scheme(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(
|
||||
module,
|
||||
SimpleNamespace(username="ops", role="admin", is_superuser=False),
|
||||
).put(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"expected_sensor_location": ["J1", "J2"],
|
||||
"sensor_location": ["J1", "J3"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["sensor_number"] == 2
|
||||
assert response.json()["sensor_location"] == ["J1", "J3"]
|
||||
|
||||
|
||||
def test_update_maps_concurrent_change_to_409(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
|
||||
def conflict(*args, **kwargs):
|
||||
raise ConflictError("方案已被其他用户修改，请重新加载")
|
||||
|
||||
monkeypatch.setattr(module, "update_sensor_placement_scheme", conflict)
|
||||
response = _client(module).put(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"expected_sensor_location": ["J1", "J2"],
|
||||
"sensor_location": ["J1", "J3"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
assert "重新加载" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_update_rejects_duplicate_nodes_before_service(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).put(
|
||||
"/api/v1/sensor-placement-schemes/7",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"expected_sensor_location": ["J1", "J2"],
|
||||
"sensor_location": ["J1", "J1"],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
def test_export_returns_xlsx_download(monkeypatch):
|
||||
module = _load_module(monkeypatch)
|
||||
response = _client(module).post(
|
||||
"/api/v1/sensor-placement-schemes/7/exports/excel",
|
||||
params={"network": "tjwater"},
|
||||
json={
|
||||
"sensor_location": ["J1", "J2"],
|
||||
"adjustment_status": {"J1": "original", "J2": "replaced"},
|
||||
},
|
||||
json={"sensor_locations": ["J1"], "adjustment_status": {}},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
@@ -426,4 +126,3 @@ def test_export_returns_xlsx_download(monkeypatch):
|
||||
assert response.headers["content-type"].startswith(
|
||||
"application/vnd.openxmlformats-officedocument"
|
||||
)
|
||||
assert "filename*=UTF-8" in response.headers["content-disposition"]
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
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.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
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.getenv("RUN_DB_INTEGRATION") != "1",
|
||||
reason="set RUN_DB_INTEGRATION=1 to test configured PostgreSQL databases",
|
||||
)
|
||||
|
||||
PROJECT = os.getenv("DB_INTEGRATION_PROJECT", "tjwater_next")
|
||||
|
||||
|
||||
def _read_business_database(_: int) -> str:
|
||||
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute("select current_database()")
|
||||
return str(cur.fetchone()["current_database"])
|
||||
|
||||
|
||||
def _read_timeseries_database(_: int) -> str:
|
||||
with timescale_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute("select current_database()")
|
||||
return str(cur.fetchone()["current_database"])
|
||||
|
||||
|
||||
def test_business_pool_handles_concurrent_borrows() -> None:
|
||||
with ThreadPoolExecutor(max_workers=16) as executor:
|
||||
names = list(executor.map(_read_business_database, range(64)))
|
||||
|
||||
assert names == [PROJECT] * 64
|
||||
|
||||
|
||||
def test_timeseries_pool_handles_concurrent_borrows() -> None:
|
||||
with ThreadPoolExecutor(max_workers=16) as executor:
|
||||
names = list(executor.map(_read_timeseries_database, range(64)))
|
||||
|
||||
assert names == [PROJECT] * 64
|
||||
|
||||
|
||||
def test_nested_wndb_writes_roll_back_as_one_transaction() -> None:
|
||||
with pytest.raises(RuntimeError, match="force rollback"):
|
||||
with project_transaction(PROJECT) as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"create temporary table wndb_pool_rollback_probe (value integer) on commit drop"
|
||||
)
|
||||
cur.execute("insert into wndb_pool_rollback_probe values (1)")
|
||||
with project_connection(PROJECT) as nested:
|
||||
assert nested is conn
|
||||
raise RuntimeError("force rollback")
|
||||
|
||||
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute("select to_regclass('pg_temp.wndb_pool_rollback_probe')")
|
||||
assert cur.fetchone()["to_regclass"] is None
|
||||
|
||||
|
||||
def test_analysis_run_lifecycle_uses_one_execution_id() -> None:
|
||||
run_id = None
|
||||
with pytest.raises(RuntimeError, match="force rollback"):
|
||||
with project_transaction(PROJECT) as conn:
|
||||
run_id = create_analysis_run(
|
||||
PROJECT,
|
||||
"integration-lifecycle-probe",
|
||||
"integration_test",
|
||||
"pytest",
|
||||
"2026-08-24T00:00:00Z",
|
||||
{"temporary": True},
|
||||
)
|
||||
update_analysis_run(
|
||||
PROJECT,
|
||||
run_id,
|
||||
status="completed",
|
||||
username="pytest",
|
||||
scheme_detail={"temporary": True},
|
||||
)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"select status from analysis.runs where run_id = %s", (run_id,)
|
||||
)
|
||||
assert cur.fetchone()["status"] == "completed"
|
||||
raise RuntimeError("force rollback")
|
||||
|
||||
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute("select count(*) as count from analysis.runs where run_id = %s", (run_id,))
|
||||
assert cur.fetchone()["count"] == 0
|
||||
|
||||
|
||||
def test_legacy_wndb_batch_treats_malicious_id_as_data() -> None:
|
||||
malicious = "integration'); DROP SCHEMA network CASCADE; --"
|
||||
command = patterns._add_pattern(
|
||||
PROJECT,
|
||||
ChangeSet({"id": malicious, "factors": [1.0]}),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="force rollback"):
|
||||
with project_transaction(PROJECT) as conn:
|
||||
write(PROJECT, command.sql)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("select count(*) as count from network.patterns where id = %s", (malicious,))
|
||||
assert cur.fetchone()["count"] == 1
|
||||
cur.execute("select to_regnamespace('network') as namespace")
|
||||
assert cur.fetchone()["namespace"] is not None
|
||||
raise RuntimeError("force rollback")
|
||||
|
||||
|
||||
def test_wndb_database_command_pattern_crud_uses_pooled_transaction() -> None:
|
||||
pattern_id = f"integration-command-{uuid4()}"
|
||||
|
||||
with pytest.raises(RuntimeError, match="force rollback"):
|
||||
with project_transaction(PROJECT) as conn:
|
||||
added = patterns.add_pattern(
|
||||
PROJECT,
|
||||
ChangeSet({"id": pattern_id, "factors": [1.0, 1.1]}),
|
||||
)
|
||||
updated = patterns.set_pattern(
|
||||
PROJECT,
|
||||
ChangeSet({"id": pattern_id, "factors": [0.8, 1.2, 1.0]}),
|
||||
)
|
||||
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"select factor from network.pattern_values "
|
||||
"where pattern_id = %s order by sequence_no",
|
||||
(pattern_id,),
|
||||
)
|
||||
assert [float(row["factor"]) for row in cur.fetchall()] == [
|
||||
0.8,
|
||||
1.2,
|
||||
1.0,
|
||||
]
|
||||
|
||||
deleted = patterns.delete_pattern(
|
||||
PROJECT,
|
||||
ChangeSet({"id": pattern_id}),
|
||||
)
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"select count(*) as count from network.patterns where id = %s",
|
||||
(pattern_id,),
|
||||
)
|
||||
assert cur.fetchone()["count"] == 0
|
||||
|
||||
assert added.operations[0]["operation"] == "add"
|
||||
assert updated.operations[0]["operation"] == "update"
|
||||
assert deleted.operations == [
|
||||
{"operation": "delete", "type": "pattern", "id": pattern_id}
|
||||
]
|
||||
raise RuntimeError("force rollback")
|
||||
|
||||
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"select count(*) as count from network.patterns where id = %s",
|
||||
(pattern_id,),
|
||||
)
|
||||
assert cur.fetchone()["count"] == 0
|
||||
|
||||
|
||||
def test_wndb_ordered_detail_tables_use_parent_scoped_primary_keys() -> None:
|
||||
expected = {
|
||||
"gis.link_vertices": "PRIMARY KEY (link_id, sequence_no)",
|
||||
"network.curve_points": "PRIMARY KEY (curve_id, sequence_no)",
|
||||
"network.demands": "PRIMARY KEY (junction_id, sequence_no)",
|
||||
"network.pattern_flow_samples": "PRIMARY KEY (pattern_id, sequence_no)",
|
||||
"network.pattern_values": "PRIMARY KEY (pattern_id, sequence_no)",
|
||||
}
|
||||
|
||||
with project_connection(PROJECT) as conn, conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"select conrelid::regclass::text as table_name, "
|
||||
"pg_get_constraintdef(oid) as definition "
|
||||
"from pg_constraint "
|
||||
"where conname = any(%s) order by 1",
|
||||
([f"{table.rsplit('.', 1)[1]}_pkey" for table in expected],),
|
||||
)
|
||||
actual = {row["table_name"]: row["definition"] for row in cur.fetchall()}
|
||||
|
||||
assert actual == expected
|
||||
|
||||
|
||||
def test_wndb_pattern_cascade_unsets_dependent_demand_atomically() -> None:
|
||||
suffix = uuid4()
|
||||
junction_id = f"integration-junction-{suffix}"
|
||||
pattern_id = f"integration-cascade-{suffix}"
|
||||
|
||||
with pytest.raises(RuntimeError, match="force rollback"):
|
||||
with project_transaction(PROJECT):
|
||||
junctions.add_junction(
|
||||
PROJECT,
|
||||
ChangeSet(
|
||||
{"id": junction_id, "x": 0.0, "y": 0.0, "elevation": 1.0}
|
||||
),
|
||||
)
|
||||
patterns.add_pattern(
|
||||
PROJECT,
|
||||
ChangeSet({"id": pattern_id, "factors": [1.0]}),
|
||||
)
|
||||
demands.set_demand(
|
||||
PROJECT,
|
||||
ChangeSet(
|
||||
{
|
||||
"junction": junction_id,
|
||||
"demands": [
|
||||
{
|
||||
"demand": 1.0,
|
||||
"pattern": pattern_id,
|
||||
"category": "integration",
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
result = delete_pattern_cascade(
|
||||
PROJECT,
|
||||
ChangeSet(g_delete_prefix | {"id": pattern_id}),
|
||||
)
|
||||
|
||||
assert patterns.get_pattern(PROJECT, pattern_id) == {}
|
||||
assert demands.get_demand(PROJECT, junction_id)["demands"] == [
|
||||
{"demand": 1.0, "pattern": None, "category": "integration"}
|
||||
]
|
||||
assert result.operations[-1] == {
|
||||
"operation": "delete",
|
||||
"type": "pattern",
|
||||
"id": pattern_id,
|
||||
}
|
||||
raise RuntimeError("force rollback")
|
||||
|
||||
assert junctions.get_junction(PROJECT, junction_id) == {}
|
||||
@@ -0,0 +1,91 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from app.infra.db.timescaledb.repositories.analysis import AnalysisResultsRepository
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def execute(self, query, params):
|
||||
self.calls.append((str(query), params))
|
||||
|
||||
async def fetchall(self):
|
||||
return []
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self):
|
||||
self.cursor_instance = _FakeCursor()
|
||||
|
||||
def cursor(self):
|
||||
return self.cursor_instance
|
||||
|
||||
|
||||
def test_prepare_simulation_rows_uses_run_timestep():
|
||||
nodes, links = AnalysisResultsRepository.prepare_simulation_rows(
|
||||
[{"node": "J1", "result": [{"pressure": 1.0}, {"pressure": 2.0}]}],
|
||||
[{"link": "P1", "result": [{"flow": 3.0}, {"flow": 4.0}]}],
|
||||
"2026-08-24T08:00:00+08:00",
|
||||
num_periods=2,
|
||||
result_timestep_seconds=900,
|
||||
)
|
||||
|
||||
assert [row["time"] for row in nodes] == [
|
||||
datetime(2026, 8, 24, 0, 0, tzinfo=timezone.utc),
|
||||
datetime(2026, 8, 24, 0, 15, tzinfo=timezone.utc),
|
||||
]
|
||||
assert nodes[0]["node_id"] == "J1"
|
||||
assert links[1]["link_id"] == "P1"
|
||||
|
||||
|
||||
def test_node_series_rejects_unknown_field():
|
||||
with pytest.raises(ValueError, match="invalid node result field"):
|
||||
asyncio.run(
|
||||
AnalysisResultsRepository.get_node_series(
|
||||
_FakeConnection(),
|
||||
uuid4(),
|
||||
"J1",
|
||||
datetime.now(timezone.utc),
|
||||
datetime.now(timezone.utc),
|
||||
"unknown",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_node_series_filters_by_run_and_node():
|
||||
conn = _FakeConnection()
|
||||
run_id = uuid4()
|
||||
start = datetime(2026, 8, 24, tzinfo=timezone.utc)
|
||||
end = datetime(2026, 8, 25, tzinfo=timezone.utc)
|
||||
|
||||
asyncio.run(
|
||||
AnalysisResultsRepository.get_node_series(
|
||||
conn, run_id, "J1", start, end, "pressure"
|
||||
)
|
||||
)
|
||||
|
||||
query, params = conn.cursor_instance.calls[0]
|
||||
assert "analysis.node_results" in query
|
||||
assert params == (run_id, "J1", start, end)
|
||||
|
||||
|
||||
def test_analysis_store_locks_run_before_empty_check():
|
||||
cursor = _FakeCursor()
|
||||
run_id = uuid4()
|
||||
|
||||
asyncio.run(AnalysisResultsRepository._lock_run(cursor, run_id))
|
||||
|
||||
query, params = cursor.calls[0]
|
||||
assert "pg_advisory_xact_lock" in query
|
||||
assert params == (run_id,)
|
||||
@@ -0,0 +1,167 @@
|
||||
import inspect
|
||||
import json
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
def test_run_simulation_exposes_explicit_valve_control():
|
||||
from app.services import simulation
|
||||
|
||||
assert "valve_control" in inspect.signature(simulation.run_simulation).parameters
|
||||
|
||||
|
||||
def test_apply_valve_control_matches_runner_semantics(monkeypatch):
|
||||
from app.services import simulation
|
||||
|
||||
updates: dict[str, dict] = {}
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"get_status",
|
||||
lambda project_name, valve_name: {
|
||||
"link": valve_name,
|
||||
"status": "OPEN",
|
||||
"setting": 1.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"set_status",
|
||||
lambda project_name, changeset: updates.update(
|
||||
{changeset.operations[0]["link"]: changeset.operations[0].copy()}
|
||||
),
|
||||
)
|
||||
|
||||
simulation._apply_valve_control(
|
||||
"demo",
|
||||
{
|
||||
"V-status": {"status": "ACTIVE"},
|
||||
"V-setting": {"setting": 2.5},
|
||||
"V-closed": {"status": "ACTIVE", "setting": 9.0, "k": 0},
|
||||
"V-k": {"status": "ACTIVE", "setting": 9.0, "k": 0.5},
|
||||
},
|
||||
)
|
||||
|
||||
assert updates["V-status"]["status"] == "ACTIVE"
|
||||
assert updates["V-setting"]["setting"] == 2.5
|
||||
assert updates["V-closed"]["status"] == "CLOSED"
|
||||
assert updates["V-k"]["setting"] == 0.1036 * pow(0.5, -3.105)
|
||||
|
||||
|
||||
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)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"get_time",
|
||||
lambda name: {
|
||||
"HYDRAULIC TIMESTEP": "00:15:00",
|
||||
"REPORT TIMESTEP": "1:00",
|
||||
"DURATION": "0:00",
|
||||
"PATTERN START": "0:00",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(simulation, "set_time", lambda name, changeset: None)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"run_project",
|
||||
lambda name: json.dumps(
|
||||
{
|
||||
"output": {
|
||||
"times": {"num_periods": 2, "report_step": 900},
|
||||
"node_results": [{"node": "J1", "result": [{}, {}]}],
|
||||
"link_results": [{"link": "P1", "result": [{}, {}]}],
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
lifecycle_calls: list[tuple] = []
|
||||
monkeypatch.setattr(simulation, "create_analysis_run", lambda **kwargs: run_id)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"update_analysis_run",
|
||||
lambda *args, **kwargs: lifecycle_calls.append((args, kwargs)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
simulation.TimescaleInternalStorage,
|
||||
"store_analysis_simulation",
|
||||
staticmethod(lambda *args, **kwargs: storage_calls.append((args, kwargs))),
|
||||
)
|
||||
|
||||
returned_run_id = simulation.run_simulation(
|
||||
name="demo",
|
||||
simulation_type="extended",
|
||||
modify_pattern_start_time="2026-07-16T00:00:00+08:00",
|
||||
modify_total_duration=900,
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="case",
|
||||
)
|
||||
|
||||
args, kwargs = storage_calls[0]
|
||||
assert args[0] == run_id
|
||||
assert args[4:] == (2, 900)
|
||||
assert kwargs["db_name"] == "demo"
|
||||
assert returned_run_id == run_id
|
||||
assert lifecycle_calls[-1][1]["status"] == "completed"
|
||||
|
||||
|
||||
def test_extended_simulation_marks_run_failed_when_result_storage_fails(monkeypatch):
|
||||
from app.services import simulation
|
||||
|
||||
run_id = uuid4()
|
||||
lifecycle_calls: list[tuple] = []
|
||||
monkeypatch.setattr(simulation, "open_project", lambda name: None)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"get_time",
|
||||
lambda name: {
|
||||
"HYDRAULIC TIMESTEP": "00:15:00",
|
||||
"REPORT TIMESTEP": "1:00",
|
||||
"DURATION": "0:00",
|
||||
"PATTERN START": "0:00",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(simulation, "set_time", lambda name, changeset: None)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"run_project",
|
||||
lambda name: json.dumps(
|
||||
{
|
||||
"output": {
|
||||
"times": {"num_periods": 1, "report_step": 900},
|
||||
"node_results": [{"node": "J1", "result": [{}]}],
|
||||
"link_results": [{"link": "P1", "result": [{}]}],
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(simulation, "create_analysis_run", lambda **kwargs: run_id)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"update_analysis_run",
|
||||
lambda *args, **kwargs: lifecycle_calls.append((args, kwargs)),
|
||||
)
|
||||
|
||||
def fail_storage(*args, **kwargs):
|
||||
raise RuntimeError("timescale write failed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
simulation.TimescaleInternalStorage,
|
||||
"store_analysis_simulation",
|
||||
staticmethod(fail_storage),
|
||||
)
|
||||
|
||||
import pytest
|
||||
|
||||
with pytest.raises(RuntimeError, match="timescale write failed"):
|
||||
simulation.run_simulation(
|
||||
name="demo",
|
||||
simulation_type="extended",
|
||||
modify_pattern_start_time="2026-07-16T00:00:00+08:00",
|
||||
modify_total_duration=900,
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="case",
|
||||
)
|
||||
|
||||
assert [call[1]["status"] for call in lifecycle_calls] == ["failed"]
|
||||
@@ -1,723 +1,101 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _load_burst_location_module():
|
||||
module_path = (
|
||||
Path(__file__).resolve().parents[2] / "app" / "services" / "burst_location.py"
|
||||
)
|
||||
|
||||
missing = object()
|
||||
previous_modules = {}
|
||||
|
||||
def install_module(name: str, module: types.ModuleType) -> None:
|
||||
previous_modules.setdefault(name, sys.modules.get(name, missing))
|
||||
sys.modules[name] = module
|
||||
|
||||
def ensure_package(name: str) -> types.ModuleType:
|
||||
module = sys.modules.get(name)
|
||||
if module is None:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = []
|
||||
install_module(name, module)
|
||||
return module
|
||||
|
||||
for package_name in [
|
||||
"app",
|
||||
"app.algorithms",
|
||||
"app.infra",
|
||||
"app.infra.db",
|
||||
"app.infra.db.timescaledb",
|
||||
"app.services",
|
||||
]:
|
||||
ensure_package(package_name)
|
||||
|
||||
time_api_module = types.ModuleType("app.services.time_api")
|
||||
time_api_module.parse_utc_time = (
|
||||
lambda value, field_name="datetime": (
|
||||
value.astimezone(timezone.utc)
|
||||
if isinstance(value, datetime) and value.tzinfo is not None
|
||||
else datetime.fromisoformat(value).astimezone(timezone.utc)
|
||||
)
|
||||
)
|
||||
time_api_module.extract_date = (
|
||||
lambda value, field_name="date": (
|
||||
value.date()
|
||||
if isinstance(value, datetime)
|
||||
else datetime.fromisoformat(value).date()
|
||||
)
|
||||
)
|
||||
time_api_module.utc_now = lambda: datetime.now(timezone.utc)
|
||||
install_module("app.services.time_api", time_api_module)
|
||||
|
||||
algorithms_module = types.ModuleType("app.algorithms.burst_location")
|
||||
algorithms_module.run_burst_location = lambda **kwargs: {}
|
||||
install_module("app.algorithms.burst_location", algorithms_module)
|
||||
|
||||
internal_queries_module = types.ModuleType(
|
||||
"app.infra.db.timescaledb.internal_queries"
|
||||
)
|
||||
|
||||
class DummyInternalQueries:
|
||||
@staticmethod
|
||||
def query_scada_by_ids_timerange(**kwargs):
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def query_scheme_simulation_by_ids_timerange(**kwargs):
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def query_realtime_simulation_by_ids_timerange(**kwargs):
|
||||
return {}
|
||||
|
||||
internal_queries_module.InternalQueries = DummyInternalQueries
|
||||
install_module(
|
||||
"app.infra.db.timescaledb.internal_queries", internal_queries_module
|
||||
)
|
||||
|
||||
scheme_management_module = types.ModuleType("app.services.scheme_management")
|
||||
scheme_management_module.query_burst_location_scheme_detail = lambda *args, **kwargs: {}
|
||||
scheme_management_module.query_burst_location_schemes = lambda *args, **kwargs: []
|
||||
scheme_management_module.query_scheme_list = lambda *args, **kwargs: []
|
||||
scheme_management_module.scheme_name_exists = lambda *args, **kwargs: False
|
||||
scheme_management_module.store_scheme_info = lambda *args, **kwargs: None
|
||||
install_module("app.services.scheme_management", scheme_management_module)
|
||||
|
||||
tjnetwork_module = types.ModuleType("app.services.tjnetwork")
|
||||
tjnetwork_module.dump_inp = lambda *args, **kwargs: None
|
||||
tjnetwork_module.get_all_scada_info = lambda *args, **kwargs: []
|
||||
install_module("app.services.tjnetwork", tjnetwork_module)
|
||||
|
||||
module_name = "tests_burst_location_under_test"
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec and spec.loader
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
finally:
|
||||
for name, previous in reversed(previous_modules.items()):
|
||||
if previous is missing:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = previous
|
||||
return module
|
||||
from app.services import burst_location
|
||||
|
||||
|
||||
def test_run_burst_location_uses_single_timerange_with_burst_source_split(monkeypatch, tmp_path):
|
||||
module = _load_burst_location_module()
|
||||
START = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
END = datetime(2026, 8, 2, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def test_analysis_simulation_query_is_keyed_by_run_id(monkeypatch):
|
||||
run_id = uuid4()
|
||||
captured = {}
|
||||
scheme_calls = []
|
||||
realtime_calls = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "pressure-query",
|
||||
},
|
||||
{
|
||||
"type": "pipe_flow",
|
||||
"associated_element_id": "P1",
|
||||
"api_query_id": "pipe-flow-query",
|
||||
},
|
||||
{
|
||||
"type": "demand",
|
||||
"associated_element_id": "J2",
|
||||
"api_query_id": "demand-query",
|
||||
},
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
|
||||
def fake_run_burst_location(**kwargs):
|
||||
def fake_query(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return {
|
||||
"located_pipe": "Pipe-001",
|
||||
"simulation_times": 3,
|
||||
"similarity_mode": "combined",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(module, "run_burst_location", fake_run_burst_location)
|
||||
|
||||
def _build_series(start_time: str, values: list[float]) -> list[dict]:
|
||||
base_time = datetime.fromisoformat(start_time)
|
||||
return [
|
||||
{
|
||||
"time": (base_time + timedelta(minutes=15 * index)).isoformat(),
|
||||
"value": value,
|
||||
}
|
||||
for index, value in enumerate(values)
|
||||
]
|
||||
|
||||
def fake_scheme_query(**kwargs):
|
||||
scheme_calls.append(kwargs)
|
||||
start_hour = datetime.fromisoformat(kwargs["start_time"]).astimezone(
|
||||
timezone(timedelta(hours=8))
|
||||
).hour
|
||||
if kwargs["element_type"] == "node" and kwargs["field"] == "pressure":
|
||||
values = [12.0, 14.0, 16.0, 18.0] if start_hour == 8 else [8.0, 10.0, 12.0, 14.0]
|
||||
return {"J1": _build_series(kwargs["start_time"], values)}
|
||||
if kwargs["element_type"] == "link" and kwargs["field"] == "flow":
|
||||
values = [5.0, 7.0, 9.0, 11.0] if start_hour == 8 else [2.0, 4.0, 6.0, 8.0]
|
||||
return {"P1": _build_series(kwargs["start_time"], values)}
|
||||
if kwargs["element_type"] == "node" and kwargs["field"] == "actual_demand":
|
||||
values = [3.0, 5.0, 7.0, 9.0] if start_hour == 8 else [1.0, 3.0, 5.0, 7.0]
|
||||
return {"J2": _build_series(kwargs["start_time"], values)}
|
||||
raise AssertionError(f"Unexpected scheme query: {kwargs}")
|
||||
|
||||
def fake_realtime_query(**kwargs):
|
||||
realtime_calls.append(kwargs)
|
||||
if kwargs["element_type"] == "node" and kwargs["field"] == "pressure":
|
||||
return {"J1": _build_series(kwargs["start_time"], [8.0, 10.0, 12.0, 14.0])}
|
||||
if kwargs["element_type"] == "link" and kwargs["field"] == "flow":
|
||||
return {"P1": _build_series(kwargs["start_time"], [2.0, 4.0, 6.0, 8.0])}
|
||||
if kwargs["element_type"] == "node" and kwargs["field"] == "actual_demand":
|
||||
return {"J2": _build_series(kwargs["start_time"], [1.0, 3.0, 5.0, 7.0])}
|
||||
raise AssertionError(f"Unexpected realtime query: {kwargs}")
|
||||
return {"J1": []}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scheme_simulation_by_ids_timerange",
|
||||
staticmethod(fake_scheme_query),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_realtime_simulation_by_ids_timerange",
|
||||
staticmethod(fake_realtime_query),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"query_scheme_list",
|
||||
lambda name, scheme_type=None: [
|
||||
(
|
||||
1,
|
||||
"BurstSchemeA",
|
||||
"burst_analysis",
|
||||
"testuser",
|
||||
None,
|
||||
None,
|
||||
{"burst_ID": ["Pipe-009", "Pipe-010"]},
|
||||
)
|
||||
],
|
||||
burst_location.InternalQueries,
|
||||
"query_analysis_simulation_by_ids_timerange",
|
||||
staticmethod(fake_query),
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="simulation",
|
||||
simulation_scheme_name="BurstSchemeA",
|
||||
simulation_scheme_type="burst_analysis",
|
||||
burst_leakage=10.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
use_scada_flow=True,
|
||||
result = burst_location._query_simulation_values(
|
||||
network="demo",
|
||||
element_ids=["J1"],
|
||||
element_type="node",
|
||||
field="pressure",
|
||||
start_dt=START,
|
||||
end_dt=END,
|
||||
simulation_source="analysis",
|
||||
simulation_run_id=run_id,
|
||||
)
|
||||
|
||||
assert result["observed_source"] == "simulation_scheme_burst_realtime_normal_timerange"
|
||||
assert result["simulation_scheme"] == {
|
||||
"name": "BurstSchemeA",
|
||||
"type": "burst_analysis",
|
||||
"burst_ids": ["Pipe-009", "Pipe-010"],
|
||||
}
|
||||
assert result["pressure_samples"] == {"burst": 4, "normal": 4}
|
||||
assert result["flow_samples"] == {"burst": 4, "normal": 4}
|
||||
assert captured["visualize_partition"] is False
|
||||
assert list(captured["burst_pressure"].index) == ["J1"]
|
||||
assert captured["burst_pressure"]["J1"] == pytest.approx(15.0)
|
||||
assert captured["normal_pressure"]["J1"] == pytest.approx(11.0)
|
||||
assert captured["burst_flow"]["J2"] == pytest.approx(6.0)
|
||||
assert captured["burst_flow"]["P1"] == pytest.approx(8.0)
|
||||
assert captured["normal_flow"]["J2"] == pytest.approx(4.0)
|
||||
assert captured["normal_flow"]["P1"] == pytest.approx(5.0)
|
||||
assert all(call["scheme_name"] == "BurstSchemeA" for call in scheme_calls)
|
||||
assert len(scheme_calls) == 3
|
||||
assert any(call["element_type"] == "node" and call["field"] == "pressure" for call in scheme_calls)
|
||||
assert any(call["element_type"] == "link" and call["field"] == "flow" for call in scheme_calls)
|
||||
assert any(call["element_type"] == "node" and call["field"] == "actual_demand" for call in scheme_calls)
|
||||
assert len(realtime_calls) == 3
|
||||
assert any(call["element_type"] == "node" and call["field"] == "pressure" for call in realtime_calls)
|
||||
assert any(call["element_type"] == "link" and call["field"] == "flow" for call in realtime_calls)
|
||||
assert any(call["element_type"] == "node" and call["field"] == "actual_demand" for call in realtime_calls)
|
||||
assert {call["start_time"] for call in scheme_calls + realtime_calls} == {
|
||||
"2025-01-01T00:00:00+00:00"
|
||||
}
|
||||
assert {call["end_time"] for call in scheme_calls + realtime_calls} == {
|
||||
"2025-01-01T01:00:00+00:00"
|
||||
}
|
||||
assert result["scada_window"] == {
|
||||
"burst_start": "2025-01-01T00:00:00+00:00",
|
||||
"burst_end": "2025-01-01T01:00:00+00:00",
|
||||
"normal_start": "2025-01-01T00:00:00+00:00",
|
||||
"normal_end": "2025-01-01T01:00:00+00:00",
|
||||
}
|
||||
assert result == {"J1": []}
|
||||
assert captured["run_id"] == run_id
|
||||
assert "scheme_name" not in captured
|
||||
assert "scheme_type" not in captured
|
||||
|
||||
|
||||
def test_run_burst_location_requires_simulation_scheme_name(monkeypatch, tmp_path):
|
||||
module = _load_burst_location_module()
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "pressure-query",
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
monkeypatch.setattr(module, "run_burst_location", lambda **kwargs: {})
|
||||
|
||||
with pytest.raises(ValueError, match="simulation_scheme_name"):
|
||||
module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="simulation",
|
||||
burst_leakage=1.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
def test_analysis_simulation_query_requires_run_id():
|
||||
with pytest.raises(ValueError, match="simulation_run_id"):
|
||||
burst_location._query_simulation_values(
|
||||
network="demo",
|
||||
element_ids=["J1"],
|
||||
element_type="node",
|
||||
field="pressure",
|
||||
start_dt=START,
|
||||
end_dt=END,
|
||||
simulation_source="analysis",
|
||||
simulation_run_id=None,
|
||||
)
|
||||
|
||||
|
||||
def test_build_observed_series_from_simulation_normalizes_result_ids(monkeypatch):
|
||||
module = _load_burst_location_module()
|
||||
query_calls = []
|
||||
|
||||
def test_scada_mapping_uses_canonical_asset_fields(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
burst_location,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": " 100026 ",
|
||||
"api_query_id": " pressure-query ",
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
def fake_scheme_query(**kwargs):
|
||||
query_calls.append(kwargs)
|
||||
return {
|
||||
100026: [
|
||||
{"time": kwargs["start_time"], "value": 10.0},
|
||||
{"time": kwargs["end_time"], "value": 14.0},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scheme_simulation_by_ids_timerange",
|
||||
staticmethod(fake_scheme_query),
|
||||
)
|
||||
|
||||
series, sample_count = module._build_observed_series_from_simulation(
|
||||
network="tjwater",
|
||||
sensor_ids=["100026"],
|
||||
start_dt=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
|
||||
end_dt=datetime(2025, 1, 1, 1, 0, 0, tzinfo=timezone.utc),
|
||||
data_type="pressure",
|
||||
series_name="burst_pressure",
|
||||
simulation_source="scheme",
|
||||
simulation_scheme_name="BurstSchemeA",
|
||||
simulation_scheme_type="burst_analysis",
|
||||
)
|
||||
|
||||
assert query_calls[0]["element_ids"] == ["100026"]
|
||||
assert sample_count == 2
|
||||
assert series["100026"] == pytest.approx(12.0)
|
||||
|
||||
|
||||
def test_build_observed_series_from_scada_uses_chinese_error_label(monkeypatch):
|
||||
module = _load_burst_location_module()
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "100026",
|
||||
"api_query_id": "pressure-query",
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(lambda **kwargs: {"pressure-query": []}),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
module._build_observed_series_from_scada(
|
||||
network="tjwater",
|
||||
sensor_ids=["100026"],
|
||||
start_dt=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
|
||||
end_dt=datetime(2025, 1, 1, 1, 0, 0, tzinfo=timezone.utc),
|
||||
data_type="pressure",
|
||||
series_name="burst_pressure",
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "爆管压力数据 在时间窗内无有效数据: 100026" in message
|
||||
assert "burst_pressure" not in message
|
||||
|
||||
|
||||
def test_build_observed_series_from_scada_skips_missing_sensor_values(monkeypatch):
|
||||
module = _load_burst_location_module()
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{"type": "pressure", "associated_element_id": "J1", "api_query_id": "q1"},
|
||||
{"type": "pressure", "associated_element_id": "J2", "api_query_id": "q2"},
|
||||
{"type": "pressure", "associated_element_id": "J3", "api_query_id": "q3"},
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(
|
||||
lambda **kwargs: {
|
||||
"q1": [
|
||||
{"time": kwargs["start_time"], "value": 10.0},
|
||||
{"time": kwargs["end_time"], "value": 12.0},
|
||||
],
|
||||
"q2": [],
|
||||
"q3": [
|
||||
{"time": kwargs["start_time"], "value": None},
|
||||
{"time": kwargs["end_time"], "value": 18.0},
|
||||
],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
series, sample_count = module._build_observed_series_from_scada(
|
||||
network="tjwater",
|
||||
sensor_ids=["J1", "J2", "J3"],
|
||||
start_dt=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
|
||||
end_dt=datetime(2025, 1, 1, 1, 0, 0, tzinfo=timezone.utc),
|
||||
data_type="pressure",
|
||||
series_name="burst_pressure",
|
||||
)
|
||||
|
||||
assert list(series.index) == ["J1", "J3"]
|
||||
assert series["J1"] == pytest.approx(11.0)
|
||||
assert series["J3"] == pytest.approx(18.0)
|
||||
assert sample_count == 1
|
||||
|
||||
|
||||
def test_run_burst_location_monitoring_uses_scada_for_burst_and_normal(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
module = _load_burst_location_module()
|
||||
captured = {}
|
||||
scada_calls = []
|
||||
realtime_calls = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "pressure-query",
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"run_burst_location",
|
||||
lambda **kwargs: captured.update(kwargs) or {"located_pipe": "Pipe-001"},
|
||||
)
|
||||
|
||||
def fake_scada_query(**kwargs):
|
||||
scada_calls.append(kwargs)
|
||||
start_hour = datetime.fromisoformat(kwargs["start_time"]).astimezone(
|
||||
timezone(timedelta(hours=8))
|
||||
).hour
|
||||
values = [20.0, 22.0] if start_hour == 8 else [10.0, 12.0]
|
||||
return {
|
||||
"pressure-query": [
|
||||
{"time": kwargs["start_time"], "value": values[0]},
|
||||
{"time": kwargs["end_time"], "value": values[1]},
|
||||
]
|
||||
}
|
||||
|
||||
def fake_realtime_query(**kwargs):
|
||||
realtime_calls.append(kwargs)
|
||||
return {
|
||||
"J1": [
|
||||
{"time": kwargs["start_time"], "value": 10.0},
|
||||
{"time": kwargs["end_time"], "value": 12.0},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(fake_scada_query),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_realtime_simulation_by_ids_timerange",
|
||||
staticmethod(fake_realtime_query),
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="monitoring",
|
||||
burst_leakage=1.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_normal_start=datetime(2025, 1, 1, 7, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_normal_end=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
)
|
||||
|
||||
assert result["observed_source"] == "scada_burst_scada_normal_timerange"
|
||||
assert len(scada_calls) == 2
|
||||
assert len(realtime_calls) == 0
|
||||
assert captured["burst_pressure"]["J1"] == pytest.approx(21.0)
|
||||
assert captured["normal_pressure"]["J1"] == pytest.approx(11.0)
|
||||
assert result["scada_window"] == {
|
||||
"burst_start": "2025-01-01T00:00:00+00:00",
|
||||
"burst_end": "2025-01-01T01:00:00+00:00",
|
||||
"normal_start": "2024-12-31T23:00:00+00:00",
|
||||
"normal_end": "2025-01-01T00:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_run_burst_location_monitoring_defaults_normal_window_to_previous_day(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
module = _load_burst_location_module()
|
||||
captured = {}
|
||||
scada_calls = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "pressure-query",
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"run_burst_location",
|
||||
lambda **kwargs: captured.update(kwargs) or {"located_pipe": "Pipe-001"},
|
||||
)
|
||||
|
||||
def fake_scada_query(**kwargs):
|
||||
scada_calls.append(kwargs)
|
||||
start_time = datetime.fromisoformat(kwargs["start_time"])
|
||||
values = (
|
||||
[20.0, 22.0]
|
||||
if start_time.date().isoformat() == "2025-01-01"
|
||||
else [10.0, 12.0]
|
||||
)
|
||||
return {
|
||||
"pressure-query": [
|
||||
{"time": kwargs["start_time"], "value": values[0]},
|
||||
{"time": kwargs["end_time"], "value": values[1]},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(fake_scada_query),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_realtime_simulation_by_ids_timerange",
|
||||
staticmethod(lambda **kwargs: pytest.fail("monitoring mode must not query realtime simulation")),
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="monitoring",
|
||||
burst_leakage=1.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
)
|
||||
|
||||
assert result["observed_source"] == "scada_burst_scada_normal_timerange"
|
||||
assert len(scada_calls) == 2
|
||||
assert datetime.fromisoformat(scada_calls[1]["start_time"]) == (
|
||||
datetime.fromisoformat(scada_calls[0]["start_time"]) - timedelta(days=1)
|
||||
)
|
||||
assert datetime.fromisoformat(scada_calls[1]["end_time"]) == (
|
||||
datetime.fromisoformat(scada_calls[0]["end_time"]) - timedelta(days=1)
|
||||
)
|
||||
assert captured["burst_pressure"]["J1"] == pytest.approx(21.0)
|
||||
assert captured["normal_pressure"]["J1"] == pytest.approx(11.0)
|
||||
assert result["scada_window"] == {
|
||||
"burst_start": "2025-01-01T00:00:00+00:00",
|
||||
"burst_end": "2025-01-01T01:00:00+00:00",
|
||||
"normal_start": "2024-12-31T00:00:00+00:00",
|
||||
"normal_end": "2024-12-31T01:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_run_burst_location_monitoring_flow_uses_previous_day_normal_window(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
module = _load_burst_location_module()
|
||||
captured = {}
|
||||
scada_calls = []
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"api_query_id": "pressure-query",
|
||||
"device_id": "pressure-1",
|
||||
"device_type": "pressure",
|
||||
"node_id": "J1",
|
||||
"link_id": None,
|
||||
"api_query_id": "q-pressure",
|
||||
},
|
||||
{
|
||||
"type": "pipe_flow",
|
||||
"associated_element_id": "P1",
|
||||
"api_query_id": "flow-query",
|
||||
"device_id": "flow-1",
|
||||
"device_type": "pipe_flow",
|
||||
"node_id": None,
|
||||
"link_id": "P1",
|
||||
"api_query_id": "q-flow",
|
||||
},
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
|
||||
assert burst_location._build_scada_mapping("demo", "pressure") == {
|
||||
"J1": "q-pressure"
|
||||
}
|
||||
assert burst_location._build_scada_mapping("demo", "flow") == {
|
||||
"P1": "q-flow"
|
||||
}
|
||||
|
||||
|
||||
def test_burst_ids_are_read_from_analysis_run(monkeypatch):
|
||||
run_id = uuid4()
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"run_burst_location",
|
||||
lambda **kwargs: captured.update(kwargs) or {"located_pipe": "Pipe-001"},
|
||||
burst_location,
|
||||
"get_analysis_run",
|
||||
lambda network, value: {
|
||||
"run_id": value,
|
||||
"parameters": {"burst_ID": ["P1", "P2"]},
|
||||
},
|
||||
)
|
||||
|
||||
def fake_scada_query(**kwargs):
|
||||
scada_calls.append(kwargs)
|
||||
is_burst_day = (
|
||||
datetime.fromisoformat(kwargs["start_time"]).date().isoformat()
|
||||
== "2025-01-01"
|
||||
)
|
||||
if kwargs["device_ids"] == ["pressure-query"]:
|
||||
values = [20.0, 22.0] if is_burst_day else [10.0, 12.0]
|
||||
query_id = "pressure-query"
|
||||
else:
|
||||
values = [7.0, 9.0] if is_burst_day else [3.0, 5.0]
|
||||
query_id = "flow-query"
|
||||
return {
|
||||
query_id: [
|
||||
{"time": kwargs["start_time"], "value": values[0]},
|
||||
{"time": kwargs["end_time"], "value": values[1]},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(fake_scada_query),
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="monitoring",
|
||||
burst_leakage=1.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
use_scada_flow=True,
|
||||
)
|
||||
|
||||
assert result["observed_source"] == "scada_burst_scada_normal_timerange"
|
||||
assert len(scada_calls) == 4
|
||||
for burst_call, normal_call in [
|
||||
(scada_calls[0], scada_calls[1]),
|
||||
(scada_calls[2], scada_calls[3]),
|
||||
]:
|
||||
assert datetime.fromisoformat(normal_call["start_time"]) == (
|
||||
datetime.fromisoformat(burst_call["start_time"]) - timedelta(days=1)
|
||||
)
|
||||
assert datetime.fromisoformat(normal_call["end_time"]) == (
|
||||
datetime.fromisoformat(burst_call["end_time"]) - timedelta(days=1)
|
||||
)
|
||||
assert captured["burst_pressure"]["J1"] == pytest.approx(21.0)
|
||||
assert captured["normal_pressure"]["J1"] == pytest.approx(11.0)
|
||||
assert captured["burst_flow"]["P1"] == pytest.approx(8.0)
|
||||
assert captured["normal_flow"]["P1"] == pytest.approx(4.0)
|
||||
|
||||
|
||||
def test_run_burst_location_monitoring_aligns_partial_scada_data(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
module = _load_burst_location_module()
|
||||
captured = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_all_scada_info",
|
||||
lambda network: [
|
||||
{"type": "pressure", "associated_element_id": "J1", "api_query_id": "q1"},
|
||||
{"type": "pressure", "associated_element_id": "J2", "api_query_id": "q2"},
|
||||
{"type": "pressure", "associated_element_id": "J3", "api_query_id": "q3"},
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(module, "_prepare_burst_inp", lambda network: str(tmp_path / "fake.inp"))
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"run_burst_location",
|
||||
lambda **kwargs: captured.update(kwargs) or {"located_pipe": "Pipe-001"},
|
||||
)
|
||||
|
||||
def fake_scada_query(**kwargs):
|
||||
start_hour = datetime.fromisoformat(kwargs["start_time"]).astimezone(
|
||||
timezone(timedelta(hours=8))
|
||||
).hour
|
||||
if start_hour == 8:
|
||||
return {
|
||||
"q1": [{"time": kwargs["start_time"], "value": 20.0}],
|
||||
"q2": [{"time": kwargs["start_time"], "value": 30.0}],
|
||||
"q3": [],
|
||||
}
|
||||
return {
|
||||
"q1": [{"time": kwargs["start_time"], "value": 10.0}],
|
||||
"q2": [],
|
||||
"q3": [{"time": kwargs["start_time"], "value": 12.0}],
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
module.InternalQueries,
|
||||
"query_scada_by_ids_timerange",
|
||||
staticmethod(fake_scada_query),
|
||||
)
|
||||
|
||||
result = module.run_burst_location_by_network(
|
||||
network="tjwater",
|
||||
username="testuser",
|
||||
data_source="monitoring",
|
||||
burst_leakage=1.0,
|
||||
scada_burst_start=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_burst_end=datetime(2025, 1, 1, 9, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_normal_start=datetime(2025, 1, 1, 7, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
scada_normal_end=datetime(2025, 1, 1, 8, 0, 0, tzinfo=timezone(timedelta(hours=8))),
|
||||
)
|
||||
|
||||
assert result["pressure_scada_ids"] == ["J1"]
|
||||
assert captured["pressure_scada_ids"] == ["J1"]
|
||||
assert list(captured["burst_pressure"].index) == ["J1"]
|
||||
assert list(captured["normal_pressure"].index) == ["J1"]
|
||||
assert captured["burst_pressure"]["J1"] == pytest.approx(20.0)
|
||||
assert captured["normal_pressure"]["J1"] == pytest.approx(10.0)
|
||||
assert burst_location._get_simulation_run_burst_ids(
|
||||
network="demo", run_id=run_id
|
||||
) == ["P1", "P2"]
|
||||
|
||||
@@ -19,15 +19,16 @@ class _FakeCursor:
|
||||
async def fetchall(self):
|
||||
return [
|
||||
{
|
||||
"id": " 25470001 ",
|
||||
"type": " PRESSURE ",
|
||||
"associated_element_id": " J1 ",
|
||||
"device_id": " 25470001 ",
|
||||
"device_type": " PRESSURE ",
|
||||
"node_id": " J1 ",
|
||||
"link_id": None,
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
"reliability": "0.95",
|
||||
"x_coor": "117.1",
|
||||
"y_coor": "32.9",
|
||||
"reliability": "95",
|
||||
"x": "117.1",
|
||||
"y": "32.9",
|
||||
}
|
||||
]
|
||||
|
||||
@@ -47,16 +48,18 @@ def test_get_scadas_normalizes_id_and_type():
|
||||
|
||||
assert result == [
|
||||
{
|
||||
"id": "25470001",
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"device_id": "25470001",
|
||||
"device_type": "pressure",
|
||||
"node_id": "J1",
|
||||
"link_id": None,
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
"reliability": 0.95,
|
||||
"reliability": 95,
|
||||
"x": 117.1,
|
||||
"y": 32.9,
|
||||
}
|
||||
]
|
||||
assert "associated_element_id" in conn.cursor_instance.query
|
||||
assert "FROM public.scada_info" in conn.cursor_instance.query
|
||||
assert "node_id" in conn.cursor_instance.query
|
||||
assert "link_id" in conn.cursor_instance.query
|
||||
assert "FROM gis.scada_devices" in conn.cursor_instance.query
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
import asyncio
|
||||
from uuid import uuid4
|
||||
|
||||
from app.infra.db.postgresql.analysis import AnalysisRepository
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def execute(self, query, params):
|
||||
self.calls.append((str(query), params))
|
||||
|
||||
async def fetchall(self):
|
||||
return []
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self):
|
||||
self.cursor_instance = _FakeCursor()
|
||||
|
||||
def cursor(self):
|
||||
return self.cursor_instance
|
||||
|
||||
|
||||
def test_list_results_uses_typed_comparison_when_result_type_is_present():
|
||||
conn = _FakeConnection()
|
||||
run_id = uuid4()
|
||||
|
||||
asyncio.run(
|
||||
AnalysisRepository.list_results(conn, run_id, "leakage_identification")
|
||||
)
|
||||
|
||||
query, params = conn.cursor_instance.calls[0]
|
||||
assert "result_type = %s" in query
|
||||
assert "%s IS NULL" not in query
|
||||
assert params == (run_id, "leakage_identification")
|
||||
|
||||
|
||||
def test_list_results_omits_result_type_filter_when_not_requested():
|
||||
conn = _FakeConnection()
|
||||
run_id = uuid4()
|
||||
|
||||
asyncio.run(AnalysisRepository.list_results(conn, run_id))
|
||||
|
||||
query, params = conn.cursor_instance.calls[0]
|
||||
assert "result_type = %s" not in query
|
||||
assert params == (run_id,)
|
||||
@@ -1,15 +1,17 @@
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock
|
||||
from uuid import uuid4
|
||||
|
||||
from app.api.v1.endpoints import project_data
|
||||
from app.infra.db.timescaledb import composite_queries
|
||||
|
||||
|
||||
PROJECT_SCADA = {
|
||||
"id": "fengyang-pressure-1",
|
||||
"type": "pressure",
|
||||
"associated_element_id": "J1",
|
||||
"device_id": "fengyang-pressure-1",
|
||||
"device_type": "pressure",
|
||||
"node_id": "J1",
|
||||
"link_id": None,
|
||||
"api_query_id": "query-1",
|
||||
"transmission_mode": "realtime",
|
||||
"transmission_frequency": None,
|
||||
@@ -42,13 +44,13 @@ def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
|
||||
composite_queries.CompositeQueries.get_scada_associated_realtime_simulation_data(
|
||||
object(),
|
||||
object(),
|
||||
[PROJECT_SCADA["id"]],
|
||||
[PROJECT_SCADA["device_id"]],
|
||||
START_TIME,
|
||||
END_TIME,
|
||||
)
|
||||
)
|
||||
|
||||
assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
|
||||
assert result[PROJECT_SCADA["device_id"]][0]["scada_id"] == PROJECT_SCADA["device_id"]
|
||||
assert query_mock.await_count == 1
|
||||
assert query_mock.await_args.args[1:] == (
|
||||
START_TIME,
|
||||
@@ -58,36 +60,35 @@ def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_scheme_scada_simulation_uses_current_project_metadata(monkeypatch):
|
||||
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}])
|
||||
monkeypatch.setattr(
|
||||
composite_queries.SchemeRepository,
|
||||
"get_node_field_by_scheme_and_time_range",
|
||||
composite_queries.AnalysisResultsRepository,
|
||||
"get_node_series",
|
||||
query_mock,
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
composite_queries.CompositeQueries.get_scada_associated_scheme_simulation_data(
|
||||
composite_queries.CompositeQueries.get_scada_associated_analysis_simulation_data(
|
||||
object(),
|
||||
object(),
|
||||
[PROJECT_SCADA["id"]],
|
||||
[PROJECT_SCADA["device_id"]],
|
||||
START_TIME,
|
||||
END_TIME,
|
||||
"baseline",
|
||||
"scheme-1",
|
||||
uuid4(),
|
||||
)
|
||||
)
|
||||
|
||||
assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
|
||||
assert query_mock.await_args.args[5:] == ("J1", "pressure")
|
||||
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")
|
||||
|
||||
|
||||
def test_element_scada_query_uses_current_project_metadata(monkeypatch):
|
||||
_patch_project_scadas(monkeypatch)
|
||||
query_mock = AsyncMock(
|
||||
return_value={
|
||||
PROJECT_SCADA["id"]: [{"time": START_TIME, "value": 26.5}]
|
||||
PROJECT_SCADA["device_id"]: [{"time": START_TIME, "value": 26.5}]
|
||||
}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from app.infra.db.timescaledb.repositories.realtime import RealtimeRepository
|
||||
@@ -65,3 +66,38 @@ def test_get_nodes_by_time_range_normalizes_inputs_to_utc():
|
||||
datetime(2026, 6, 1, 0, 0, tzinfo=timezone.utc),
|
||||
datetime(2026, 6, 1, 1, 0, tzinfo=timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
class _SyncTransactionConnection:
|
||||
def __init__(self):
|
||||
self.transactions = 0
|
||||
|
||||
@contextmanager
|
||||
def transaction(self):
|
||||
self.transactions += 1
|
||||
yield
|
||||
|
||||
|
||||
def test_realtime_node_and_link_replacement_share_outer_transaction(monkeypatch):
|
||||
conn = _SyncTransactionConnection()
|
||||
calls: list[str] = []
|
||||
monkeypatch.setattr(
|
||||
RealtimeRepository,
|
||||
"insert_nodes_batch_sync",
|
||||
lambda _conn, _data: calls.append("nodes"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
RealtimeRepository,
|
||||
"insert_links_batch_sync",
|
||||
lambda _conn, _data: calls.append("links"),
|
||||
)
|
||||
|
||||
RealtimeRepository.store_realtime_simulation_result_sync(
|
||||
conn,
|
||||
[{"node": "N1", "result": [{"pressure": 1.0}]}],
|
||||
[{"link": "L1", "result": [{"flow": 2.0}]}],
|
||||
"2026-06-01T00:00:00Z",
|
||||
)
|
||||
|
||||
assert conn.transactions == 1
|
||||
assert calls == ["nodes", "links"]
|
||||
|
||||
@@ -17,7 +17,7 @@ def test_clean_scada_uses_current_project_metadata(monkeypatch):
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
return_value=[{"device_id": "fengyang-pressure-1", "device_type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
@@ -66,7 +66,7 @@ def test_clean_scada_rejects_devices_missing_from_project_metadata(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(return_value=[{"id": "other-device", "type": "pressure"}]),
|
||||
AsyncMock(return_value=[{"device_id": "other-device", "device_type": "pressure"}]),
|
||||
)
|
||||
query_mock = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
@@ -94,7 +94,7 @@ def test_clean_scada_rejects_zero_database_updates(monkeypatch):
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
return_value=[{"device_id": "fengyang-pressure-1", "device_type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
@@ -141,7 +141,7 @@ def test_clean_scada_propagates_write_failures(monkeypatch):
|
||||
composite_queries.ScadaInfoRepository,
|
||||
"get_scadas",
|
||||
AsyncMock(
|
||||
return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
|
||||
return_value=[{"device_id": "fengyang-pressure-1", "device_type": "pressure"}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
|
||||
@@ -65,8 +65,8 @@ def test_update_scada_field_inserts_when_update_hits_no_rows():
|
||||
)
|
||||
|
||||
assert len(conn.cursor_instance.calls) == 2
|
||||
assert "UPDATE scada.scada_data SET" in conn.cursor_instance.calls[0][0]
|
||||
assert "INSERT INTO scada.scada_data" in conn.cursor_instance.calls[1][0]
|
||||
assert "UPDATE scada.measurements SET" in conn.cursor_instance.calls[0][0]
|
||||
assert "INSERT INTO scada.measurements" in conn.cursor_instance.calls[1][0]
|
||||
|
||||
|
||||
def test_update_scada_field_skips_insert_when_update_succeeds():
|
||||
@@ -85,4 +85,4 @@ def test_update_scada_field_skips_insert_when_update_succeeds():
|
||||
)
|
||||
|
||||
assert len(conn.cursor_instance.calls) == 1
|
||||
assert "UPDATE scada.scada_data SET" in conn.cursor_instance.calls[0][0]
|
||||
assert "UPDATE scada.measurements SET" in conn.cursor_instance.calls[0][0]
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
from app.services import scheme_management, tjnetwork
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_exc_info):
|
||||
return False
|
||||
|
||||
def execute(self, statement, params=None):
|
||||
self.calls.append((str(statement), params))
|
||||
|
||||
def fetchall(self):
|
||||
return []
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self, cursor):
|
||||
self._cursor = cursor
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *_exc_info):
|
||||
return False
|
||||
|
||||
def cursor(self):
|
||||
return self._cursor
|
||||
|
||||
|
||||
def test_query_scheme_list_pushes_scheme_type_into_sql(monkeypatch):
|
||||
cursor = _FakeCursor()
|
||||
monkeypatch.setattr(
|
||||
scheme_management,
|
||||
"get_project_pgconn_string",
|
||||
lambda db_name=None: "postgres://test",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
scheme_management.psycopg, "connect", lambda _conn_string: _FakeConnection(cursor)
|
||||
)
|
||||
|
||||
assert scheme_management.query_scheme_list("demo", scheme_type="burst_analysis") == []
|
||||
|
||||
statement, params = cursor.calls[0]
|
||||
assert "WHERE scheme_type = %s" in statement
|
||||
assert params == ("burst_analysis",)
|
||||
|
||||
|
||||
def test_get_all_schemes_filters_central_scheme_list_by_type(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_query_scheme_list(name, scheme_type=None, query_date=None):
|
||||
captured["name"] = name
|
||||
captured["scheme_type"] = scheme_type
|
||||
captured["query_date"] = query_date
|
||||
return [
|
||||
(
|
||||
7,
|
||||
"burst_case",
|
||||
"burst_analysis",
|
||||
"alice",
|
||||
"2026-01-01T00:00:00+08:00",
|
||||
"2026-01-01T01:00:00+08:00",
|
||||
{"burst_ID": ["P1"]},
|
||||
)
|
||||
]
|
||||
|
||||
monkeypatch.setattr(
|
||||
scheme_management, "query_scheme_list", fake_query_scheme_list
|
||||
)
|
||||
|
||||
result = tjnetwork.get_all_schemes("demo", scheme_type="burst_analysis")
|
||||
|
||||
assert captured == {
|
||||
"name": "demo",
|
||||
"scheme_type": "burst_analysis",
|
||||
"query_date": None,
|
||||
}
|
||||
assert result == [
|
||||
{
|
||||
"scheme_id": 7,
|
||||
"scheme_name": "burst_case",
|
||||
"scheme_type": "burst_analysis",
|
||||
"username": "alice",
|
||||
"create_time": "2026-01-01T00:00:00+08:00",
|
||||
"scheme_start_time": "2026-01-01T01:00:00+08:00",
|
||||
"scheme_detail": {"burst_ID": ["P1"]},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_query_scheme_detail_rejects_wrong_specialized_type(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
scheme_management,
|
||||
"query_burst_detection_scheme_detail",
|
||||
lambda name, scheme_name: {
|
||||
"scheme_name": scheme_name,
|
||||
"scheme_type": "burst_analysis",
|
||||
"network": name,
|
||||
},
|
||||
)
|
||||
|
||||
assert (
|
||||
scheme_management.query_scheme_detail(
|
||||
"demo",
|
||||
"same_name",
|
||||
scheme_type="burst_detection",
|
||||
)
|
||||
== {}
|
||||
)
|
||||
|
||||
|
||||
def test_query_scheme_detail_rejects_wrong_network(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
scheme_management,
|
||||
"query_burst_location_scheme_detail",
|
||||
lambda name, scheme_name: {
|
||||
"scheme_name": scheme_name,
|
||||
"scheme_type": "burst_location",
|
||||
"network": "other_network",
|
||||
},
|
||||
)
|
||||
|
||||
assert (
|
||||
scheme_management.query_scheme_detail(
|
||||
"demo",
|
||||
"same_name",
|
||||
scheme_type="burst_location",
|
||||
)
|
||||
== {}
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from app.services import scheme_management
|
||||
|
||||
|
||||
def _mock_connection(monkeypatch):
|
||||
cursor = MagicMock()
|
||||
cursor.rowcount = 1
|
||||
connection = MagicMock()
|
||||
connection.cursor.return_value.__enter__.return_value = cursor
|
||||
context = MagicMock()
|
||||
context.__enter__.return_value = connection
|
||||
monkeypatch.setattr(scheme_management, "project_connection", lambda _name: context)
|
||||
return cursor
|
||||
|
||||
|
||||
def test_repeated_analysis_names_create_distinct_run_ids(monkeypatch) -> None:
|
||||
cursor = _mock_connection(monkeypatch)
|
||||
arguments = {
|
||||
"name": "tjwater_next",
|
||||
"scheme_name": "same-window",
|
||||
"scheme_type": "burst_analysis",
|
||||
"username": "alice",
|
||||
"scheme_start_time": "2026-08-24T00:00:00Z",
|
||||
"scheme_detail": {},
|
||||
}
|
||||
|
||||
first = scheme_management.create_analysis_run(**arguments)
|
||||
second = scheme_management.create_analysis_run(**arguments)
|
||||
|
||||
assert first != second
|
||||
assert cursor.execute.call_count == 2
|
||||
assert all(
|
||||
"insert into analysis.runs" in call.args[0]
|
||||
for call in cursor.execute.call_args_list
|
||||
)
|
||||
|
||||
|
||||
def test_update_analysis_run_targets_execution_id(monkeypatch) -> None:
|
||||
cursor = _mock_connection(monkeypatch)
|
||||
run_id = scheme_management.create_analysis_run(
|
||||
name="tjwater_next",
|
||||
scheme_name="run",
|
||||
scheme_type="burst_analysis",
|
||||
username="alice",
|
||||
scheme_start_time="2026-08-24T00:00:00Z",
|
||||
scheme_detail={},
|
||||
)
|
||||
cursor.reset_mock()
|
||||
|
||||
scheme_management.update_analysis_run(
|
||||
"tjwater_next",
|
||||
run_id,
|
||||
status="completed",
|
||||
username="alice",
|
||||
scheme_detail={"window": "24h"},
|
||||
)
|
||||
|
||||
statement, params = cursor.execute.call_args.args
|
||||
assert "where run_id = %s" in statement
|
||||
assert params[-1] == run_id
|
||||
assert params[1] == "completed"
|
||||
@@ -1,219 +0,0 @@
|
||||
import inspect
|
||||
import json
|
||||
from datetime import timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from app.infra.db.timescaledb.repositories.scheme import SchemeRepository
|
||||
from app.services.time_api import parse_utc_time
|
||||
|
||||
|
||||
def test_run_simulation_exposes_explicit_valve_control():
|
||||
from app.services import simulation
|
||||
|
||||
parameters = inspect.signature(simulation.run_simulation).parameters
|
||||
|
||||
assert "valve_control" in parameters
|
||||
|
||||
|
||||
def test_apply_valve_control_matches_run_simulation_ex_semantics(monkeypatch):
|
||||
from app.services import simulation
|
||||
|
||||
updates: dict[str, dict] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"get_status",
|
||||
lambda project_name, valve_name: {
|
||||
"link": valve_name,
|
||||
"status": "OPEN",
|
||||
"setting": 1.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"set_status",
|
||||
lambda project_name, changeset: updates.update(
|
||||
{
|
||||
changeset.operations[0]["link"]: changeset.operations[0].copy()
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
simulation._apply_valve_control(
|
||||
"demo",
|
||||
{
|
||||
"V-status": {"status": "ACTIVE"},
|
||||
"V-setting": {"setting": 2.5},
|
||||
"V-closed": {"status": "ACTIVE", "setting": 9.0, "k": 0},
|
||||
"V-k": {"status": "ACTIVE", "setting": 9.0, "k": 0.5},
|
||||
},
|
||||
)
|
||||
|
||||
assert updates["V-status"] == {
|
||||
"link": "V-status",
|
||||
"status": "ACTIVE",
|
||||
"setting": 1.0,
|
||||
}
|
||||
assert updates["V-setting"] == {
|
||||
"link": "V-setting",
|
||||
"status": "OPEN",
|
||||
"setting": 2.5,
|
||||
}
|
||||
assert updates["V-closed"] == {
|
||||
"link": "V-closed",
|
||||
"status": "CLOSED",
|
||||
"setting": 9.0,
|
||||
}
|
||||
assert updates["V-k"] == {
|
||||
"link": "V-k",
|
||||
"status": "ACTIVE",
|
||||
"setting": 0.1036 * pow(0.5, -3.105),
|
||||
}
|
||||
|
||||
|
||||
def _node_result(periods: int) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"node": "J1",
|
||||
"result": [
|
||||
{"demand": index, "head": index, "pressure": index, "quality": index}
|
||||
for index in range(periods)
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _link_result(periods: int) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"link": "P1",
|
||||
"result": [
|
||||
{
|
||||
"flow": index,
|
||||
"friction": index,
|
||||
"headloss": index,
|
||||
"quality": index,
|
||||
"reaction": index,
|
||||
"setting": index,
|
||||
"status": index,
|
||||
"velocity": index,
|
||||
}
|
||||
for index in range(periods)
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_store_scheme_simulation_uses_15_minute_report_step(monkeypatch):
|
||||
inserted: dict[str, list[dict]] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_nodes_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_links_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("links", data)),
|
||||
)
|
||||
|
||||
SchemeRepository.store_scheme_simulation_result_sync(
|
||||
conn=object(),
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="five_hour_case",
|
||||
node_result_list=_node_result(21),
|
||||
link_result_list=_link_result(21),
|
||||
result_start_time="2026-07-16T00:00:00Z",
|
||||
num_periods=21,
|
||||
result_timestep_seconds=900,
|
||||
)
|
||||
|
||||
start_time = parse_utc_time("2026-07-16T00:00:00Z")
|
||||
assert len(inserted["nodes"]) == 21
|
||||
assert inserted["nodes"][0]["time"] == start_time
|
||||
assert inserted["nodes"][-1]["time"] == start_time + timedelta(hours=5)
|
||||
assert inserted["links"][-1]["time"] == start_time + timedelta(hours=5)
|
||||
|
||||
|
||||
def test_store_scheme_simulation_uses_hourly_report_step(monkeypatch):
|
||||
inserted: dict[str, list[dict]] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_nodes_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("nodes", data)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
SchemeRepository,
|
||||
"insert_links_batch_sync",
|
||||
staticmethod(lambda conn, data: inserted.setdefault("links", data)),
|
||||
)
|
||||
|
||||
SchemeRepository.store_scheme_simulation_result_sync(
|
||||
conn=object(),
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="hourly_case",
|
||||
node_result_list=_node_result(6),
|
||||
link_result_list=_link_result(6),
|
||||
result_start_time="2026-07-16T00:00:00Z",
|
||||
num_periods=6,
|
||||
result_timestep_seconds=3600,
|
||||
)
|
||||
|
||||
start_time = parse_utc_time("2026-07-16T00:00:00Z")
|
||||
assert [item["time"] for item in inserted["nodes"]] == [
|
||||
start_time + timedelta(hours=index) for index in range(6)
|
||||
]
|
||||
|
||||
|
||||
def test_run_simulation_passes_report_step_for_extended_scheme(monkeypatch):
|
||||
import app.services.simulation as simulation
|
||||
|
||||
time_updates: list[dict] = []
|
||||
storage_calls: list[tuple] = []
|
||||
|
||||
monkeypatch.setattr(simulation, "open_project", lambda name: None)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"get_time",
|
||||
lambda name: {
|
||||
"HYDRAULIC TIMESTEP": "00:15:00",
|
||||
"REPORT TIMESTEP": "1:00",
|
||||
"DURATION": "0:00",
|
||||
"PATTERN START": "0:00",
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
simulation,
|
||||
"set_time",
|
||||
lambda name, changeset: time_updates.append(changeset.operations[0]),
|
||||
)
|
||||
monkeypatch.setattr(simulation, "run_project", lambda name: json.dumps({
|
||||
"simulation_result": "successful",
|
||||
"output": {
|
||||
"times": {"num_periods": 21, "report_step": 900},
|
||||
"node_results": _node_result(21),
|
||||
"link_results": _link_result(21),
|
||||
},
|
||||
}))
|
||||
monkeypatch.setattr(
|
||||
simulation.TimescaleInternalStorage,
|
||||
"store_scheme_simulation",
|
||||
staticmethod(lambda *args, **kwargs: storage_calls.append((args, kwargs))),
|
||||
)
|
||||
|
||||
simulation.run_simulation(
|
||||
name="fengyang",
|
||||
simulation_type="extended",
|
||||
modify_pattern_start_time="2026-07-16T00:00:00+08:00",
|
||||
modify_total_duration=18000,
|
||||
scheme_type="burst_analysis",
|
||||
scheme_name="five_hour_case",
|
||||
)
|
||||
|
||||
assert time_updates[0]["DURATION"] == "05:00:00"
|
||||
assert time_updates[0]["REPORT TIMESTEP"] == "1:00"
|
||||
assert storage_calls[0][0][5] == 21
|
||||
assert storage_calls[0][0][6] == 900
|
||||
@@ -1,10 +1,11 @@
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from openpyxl import load_workbook
|
||||
|
||||
from app.native.wndb import s42_sensor_placement
|
||||
from app.infra.db.postgresql import sensor_placement as sensor_placement_repository
|
||||
from app.services import sensor_placement
|
||||
|
||||
|
||||
@@ -15,180 +16,121 @@ def _mock_project_cursor(monkeypatch):
|
||||
connection_context = MagicMock()
|
||||
connection_context.__enter__.return_value = connection
|
||||
monkeypatch.setattr(
|
||||
s42_sensor_placement,
|
||||
sensor_placement_repository,
|
||||
"project_connection",
|
||||
lambda _network: connection_context,
|
||||
)
|
||||
return cursor
|
||||
|
||||
|
||||
def test_build_workbook_contains_engineering_columns(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
sensor_placement,
|
||||
"_sensor_points",
|
||||
lambda network, locations: [
|
||||
{
|
||||
"node_id": "J1",
|
||||
"max_pipe_diameter": 400.0,
|
||||
"project_x": 13500000.0,
|
||||
"project_y": 3600000.0,
|
||||
"map_x": 13500000.0,
|
||||
"map_y": 3600000.0,
|
||||
"longitude": 121.0,
|
||||
"latitude": 31.0,
|
||||
"elevation": 4.5,
|
||||
}
|
||||
],
|
||||
)
|
||||
scheme = {
|
||||
"id": 7,
|
||||
"scheme_name": "北区测压点",
|
||||
"sensor_number": 2,
|
||||
def _run() -> dict:
|
||||
return {
|
||||
"run_id": uuid4(),
|
||||
"name": "北区测压点",
|
||||
"sensor_count": 2,
|
||||
"min_diameter": 300,
|
||||
"username": "alice",
|
||||
"create_time": datetime(2026, 7, 30, tzinfo=timezone.utc),
|
||||
"sensor_location": ["J1", "J2"],
|
||||
"created_by": "alice",
|
||||
"created_at": datetime(2026, 7, 30, tzinfo=timezone.utc),
|
||||
"status": "completed",
|
||||
"sensor_locations": ["J1", "J2"],
|
||||
}
|
||||
|
||||
|
||||
def _point() -> dict:
|
||||
return {
|
||||
"node_id": "J1",
|
||||
"max_pipe_diameter": 400.0,
|
||||
"project_x": 3038.94,
|
||||
"project_y": -34446.59,
|
||||
"map_x": 13525191.530279,
|
||||
"map_y": 3622984.760237,
|
||||
"longitude": 121.498863,
|
||||
"latitude": 30.924784,
|
||||
"elevation": 4.5,
|
||||
}
|
||||
|
||||
|
||||
def test_build_workbook_uses_analysis_run_metadata(monkeypatch):
|
||||
monkeypatch.setattr(sensor_placement, "_sensor_points", lambda *_: [_point()])
|
||||
output = sensor_placement.build_sensor_placement_workbook(
|
||||
network="tjwater",
|
||||
scheme=scheme,
|
||||
scheme=_run(),
|
||||
sensor_location=["J1"],
|
||||
adjustment_status={"J1": "replaced"},
|
||||
)
|
||||
workbook = load_workbook(output)
|
||||
|
||||
assert workbook.sheetnames == ["方案信息", "监测点清单"]
|
||||
headers = [cell.value for cell in workbook["监测点清单"][1]]
|
||||
assert headers == [
|
||||
"序号",
|
||||
"节点 ID",
|
||||
"经度",
|
||||
"纬度",
|
||||
"工程 X",
|
||||
"工程 Y",
|
||||
"地图 X",
|
||||
"地图 Y",
|
||||
"高程",
|
||||
"调整状态",
|
||||
]
|
||||
assert workbook["监测点清单"]["J2"].value == "替换"
|
||||
assert workbook["方案信息"]["B8"].value == "未保存草稿"
|
||||
|
||||
|
||||
def test_candidate_keeps_engineering_coordinates_and_transforms_map_coordinates(
|
||||
monkeypatch,
|
||||
):
|
||||
def test_candidate_transforms_map_coordinates(monkeypatch):
|
||||
row = _point().copy()
|
||||
row.pop("longitude")
|
||||
row.pop("latitude")
|
||||
monkeypatch.setattr(
|
||||
sensor_placement.wndb,
|
||||
sensor_placement.sensor_placement_repository,
|
||||
"get_sensor_placement_nodes",
|
||||
lambda network, node_ids: [
|
||||
{
|
||||
"node_id": "J1",
|
||||
"max_pipe_diameter": 400.0,
|
||||
"project_x": 3038.94,
|
||||
"project_y": -34446.59,
|
||||
"map_x": 13525191.530279,
|
||||
"map_y": 3622984.760237,
|
||||
"elevation": 4.5,
|
||||
}
|
||||
],
|
||||
lambda network, node_ids: [row],
|
||||
)
|
||||
|
||||
point = sensor_placement.get_sensor_placement_candidate("tjwater", "J1")
|
||||
|
||||
assert point["project_x"] == 3038.94
|
||||
assert point["project_y"] == -34446.59
|
||||
assert point["max_pipe_diameter"] == 400.0
|
||||
assert point["longitude"] == pytest.approx(121.498863, abs=1e-6)
|
||||
assert point["latitude"] == pytest.approx(30.924784, abs=1e-6)
|
||||
|
||||
|
||||
def test_update_validates_nodes_before_write(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
sensor_placement.wndb,
|
||||
sensor_placement.sensor_placement_repository,
|
||||
"get_sensor_placement_nodes",
|
||||
lambda network, node_ids: [],
|
||||
)
|
||||
|
||||
try:
|
||||
sensor_placement.update_sensor_placement_scheme(
|
||||
with pytest.raises(sensor_placement.SensorPlacementValidationError, match="missing"):
|
||||
sensor_placement.update_sensor_placement_run(
|
||||
"tjwater",
|
||||
7,
|
||||
expected_sensor_location=["J1"],
|
||||
sensor_location=["missing"],
|
||||
uuid4(),
|
||||
expected_sensor_locations=["J1"],
|
||||
sensor_locations=["missing"],
|
||||
)
|
||||
except sensor_placement.SensorPlacementValidationError as exc:
|
||||
assert "missing" in str(exc)
|
||||
else:
|
||||
raise AssertionError("expected invalid node to be rejected")
|
||||
|
||||
|
||||
def test_sensor_nodes_use_materialized_web_mercator_geometry(monkeypatch):
|
||||
def test_sensor_nodes_query_uses_new_network_and_gis_schemas(monkeypatch):
|
||||
cursor = _mock_project_cursor(monkeypatch)
|
||||
cursor.fetchall.return_value = []
|
||||
|
||||
s42_sensor_placement.get_sensor_placement_nodes("tjwater", ["J1"])
|
||||
sensor_placement_repository.get_sensor_placement_nodes("tjwater", ["J1"])
|
||||
|
||||
query = cursor.execute.call_args.args[0]
|
||||
assert "geo_junctions_mat" in query
|
||||
assert "ST_X(c.coord)" in query
|
||||
assert "ST_Y(c.coord)" in query
|
||||
assert "ST_X(gj.geom)" in query
|
||||
assert "ST_Y(gj.geom)" in query
|
||||
assert "MAX(diameter) AS max_pipe_diameter" in query
|
||||
assert "network.pipes" in query
|
||||
assert "network.links" in query
|
||||
assert "gis.node_geometries" in query
|
||||
assert "ST_Transform(g.geom, 3857)" in query
|
||||
assert cursor.execute.call_args.args[1] == (["J1"], ["J1"], ["J1"])
|
||||
|
||||
|
||||
def test_workbook_escapes_formula_in_scheme_metadata(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
sensor_placement,
|
||||
"_sensor_points",
|
||||
lambda network, locations: [
|
||||
{
|
||||
"node_id": "J1",
|
||||
"max_pipe_diameter": 400.0,
|
||||
"project_x": 3038.94,
|
||||
"project_y": -34446.59,
|
||||
"map_x": 13525191.53,
|
||||
"map_y": 3622984.76,
|
||||
"longitude": 121.49,
|
||||
"latitude": 30.92,
|
||||
"elevation": 4.5,
|
||||
}
|
||||
],
|
||||
)
|
||||
output = sensor_placement.build_sensor_placement_workbook(
|
||||
network="tjwater",
|
||||
scheme={
|
||||
"scheme_name": "=1+1",
|
||||
"sensor_location": ["J1"],
|
||||
"min_diameter": 300,
|
||||
"username": "alice",
|
||||
"create_time": datetime(2026, 7, 30, tzinfo=timezone.utc),
|
||||
},
|
||||
sensor_location=["J1"],
|
||||
adjustment_status={},
|
||||
)
|
||||
|
||||
workbook = load_workbook(output, data_only=False)
|
||||
assert workbook["方案信息"]["B2"].value == "'=1+1"
|
||||
assert workbook["方案信息"]["B2"].data_type == "s"
|
||||
|
||||
|
||||
def test_create_sensor_placement_returns_inserted_record(monkeypatch):
|
||||
def test_create_sensor_placement_writes_run_and_result_atomically(monkeypatch):
|
||||
cursor = _mock_project_cursor(monkeypatch)
|
||||
cursor.fetchone.return_value = {"id": 7, "sensor_location": ["J1", "J2"]}
|
||||
created_at = datetime(2026, 8, 24, tzinfo=timezone.utc)
|
||||
cursor.fetchone.return_value = {
|
||||
"run_id": uuid4(),
|
||||
"name": "北区测压点",
|
||||
"created_by": "alice",
|
||||
"created_at": created_at,
|
||||
"status": "completed",
|
||||
}
|
||||
|
||||
created = s42_sensor_placement.create_sensor_placement(
|
||||
created = sensor_placement_repository.create_sensor_placement(
|
||||
"tjwater",
|
||||
scheme_name="北区测压点",
|
||||
run_name="北区测压点",
|
||||
min_diameter=300,
|
||||
username="alice",
|
||||
sensor_location=["J1", "J2"],
|
||||
created_by="alice",
|
||||
sensor_locations=["J1", "J2"],
|
||||
)
|
||||
|
||||
assert created["id"] == 7
|
||||
query, parameters = cursor.execute.call_args.args
|
||||
assert "INSERT INTO sensor_placement" in query
|
||||
assert parameters == ("北区测压点", 2, 300, "alice", ["J1", "J2"])
|
||||
assert created["sensor_count"] == 2
|
||||
statements = [call.args[0] for call in cursor.execute.call_args_list]
|
||||
assert "INSERT INTO analysis.runs" in statements[0]
|
||||
assert "INSERT INTO analysis.results" in statements[1]
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from app.infra.db.timescaledb import sync_pool
|
||||
|
||||
|
||||
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 object()
|
||||
finally:
|
||||
self.returned += 1
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_pools():
|
||||
sync_pool._pools.clear()
|
||||
sync_pool._pool_conninfo.clear()
|
||||
sync_pool._pool_borrows.clear()
|
||||
yield
|
||||
sync_pool._pools.clear()
|
||||
sync_pool._pool_conninfo.clear()
|
||||
sync_pool._pool_borrows.clear()
|
||||
|
||||
|
||||
def test_pool_reuses_same_routed_timescale_dsn(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
first = sync_pool.get_timescale_pool("tjwater_next")
|
||||
second = sync_pool.get_timescale_pool("tjwater_next")
|
||||
|
||||
assert first is second
|
||||
|
||||
|
||||
def test_pool_rebuilds_after_routing_change(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
dsn = {"value": "host=old dbname=tjwater_next"}
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: dsn["value"],
|
||||
)
|
||||
|
||||
old = sync_pool.get_timescale_pool("tjwater_next")
|
||||
dsn["value"] = "host=new dbname=tjwater_next"
|
||||
new = sync_pool.get_timescale_pool("tjwater_next")
|
||||
|
||||
assert old.closed is True
|
||||
assert new is not old
|
||||
|
||||
|
||||
def test_connection_is_returned_to_pool(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
pool = sync_pool.get_timescale_pool("tjwater_next")
|
||||
with sync_pool.timescale_connection("tjwater_next"):
|
||||
assert pool.borrowed == 1
|
||||
assert pool.returned == 0
|
||||
assert pool.returned == 1
|
||||
|
||||
|
||||
def test_close_removes_pool(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
pool = sync_pool.get_timescale_pool("tjwater_next")
|
||||
sync_pool.close_timescale_pool("tjwater_next")
|
||||
|
||||
assert pool.closed is True
|
||||
assert "tjwater_next" not in sync_pool._pools
|
||||
|
||||
|
||||
def test_close_all_removes_every_pool(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
first = sync_pool.get_timescale_pool("first")
|
||||
second = sync_pool.get_timescale_pool("second")
|
||||
sync_pool.close_all_timescale_pools()
|
||||
|
||||
assert first.closed is True
|
||||
assert second.closed is True
|
||||
assert sync_pool._pools == {}
|
||||
|
||||
|
||||
def test_pool_cache_evicts_least_recently_used_idle_pool(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(sync_pool.settings, "PROJECT_TS_CACHE_SIZE", 2)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
first = sync_pool.get_timescale_pool("first")
|
||||
second = sync_pool.get_timescale_pool("second")
|
||||
sync_pool.get_timescale_pool("first")
|
||||
third = sync_pool.get_timescale_pool("third")
|
||||
|
||||
assert list(sync_pool._pools) == ["first", "third"]
|
||||
assert second.closed is True
|
||||
assert first.closed is False
|
||||
assert third.closed is False
|
||||
|
||||
|
||||
def test_pool_cache_does_not_evict_active_pool(monkeypatch):
|
||||
monkeypatch.setattr(sync_pool, "ConnectionPool", _FakePool)
|
||||
monkeypatch.setattr(sync_pool.settings, "PROJECT_TS_CACHE_SIZE", 1)
|
||||
monkeypatch.setattr(
|
||||
sync_pool,
|
||||
"get_project_timescale_pgconn_string",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
|
||||
with sync_pool.timescale_connection("active"):
|
||||
active = sync_pool._pools["active"]
|
||||
sync_pool.get_timescale_pool("new")
|
||||
assert active.closed is False
|
||||
assert set(sync_pool._pools) == {"active", "new"}
|
||||
|
||||
assert list(sync_pool._pools) == ["new"]
|
||||
assert active.closed is True
|
||||
@@ -0,0 +1,83 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from app.native.wndb.commands import executor
|
||||
from app.native.wndb.core.database import ChangeSet
|
||||
|
||||
|
||||
def test_batch_commits_before_materialized_view_refresh(monkeypatch) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
@contextmanager
|
||||
def fake_transaction(_name: str):
|
||||
events.append("transaction-enter")
|
||||
yield object()
|
||||
events.append("transaction-exit")
|
||||
|
||||
monkeypatch.setattr(executor, "project_transaction", fake_transaction)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"expand_command",
|
||||
lambda _name, change_set: change_set,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"_execute_update_command",
|
||||
lambda _name, _change_set: events.append("write") or ChangeSet(),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"refresh_materialized_views",
|
||||
lambda _name: events.append("refresh"),
|
||||
)
|
||||
|
||||
executor.execute_batch_commands(
|
||||
"project_a",
|
||||
ChangeSet({"operation": "update", "type": "junction", "id": "J1"}),
|
||||
)
|
||||
|
||||
assert events == [
|
||||
"transaction-enter",
|
||||
"write",
|
||||
"transaction-exit",
|
||||
"refresh",
|
||||
]
|
||||
|
||||
|
||||
def test_failed_batch_does_not_refresh_materialized_views(monkeypatch) -> None:
|
||||
events: list[str] = []
|
||||
|
||||
@contextmanager
|
||||
def fake_transaction(_name: str):
|
||||
events.append("transaction-enter")
|
||||
try:
|
||||
yield object()
|
||||
finally:
|
||||
events.append("transaction-exit")
|
||||
|
||||
monkeypatch.setattr(executor, "project_transaction", fake_transaction)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"expand_command",
|
||||
lambda _name, change_set: change_set,
|
||||
)
|
||||
|
||||
def fail_write(_name, _change_set):
|
||||
events.append("write")
|
||||
raise RuntimeError("write failed")
|
||||
|
||||
monkeypatch.setattr(executor, "_execute_update_command", fail_write)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"refresh_materialized_views",
|
||||
lambda _name: events.append("refresh"),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="write failed"):
|
||||
executor.execute_batch_commands(
|
||||
"project_a",
|
||||
ChangeSet({"operation": "update", "type": "junction", "id": "J1"}),
|
||||
)
|
||||
|
||||
assert events == ["transaction-enter", "write", "transaction-exit"]
|
||||
+197
-114
@@ -1,145 +1,228 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from app.native.wndb import connection
|
||||
from app.native.wndb import database
|
||||
from app.native.wndb import project
|
||||
from app.native.wndb.core import connection
|
||||
|
||||
|
||||
class _FakeCursor:
|
||||
def __init__(self, connection):
|
||||
self.connection = connection
|
||||
class _FakePool:
|
||||
def __init__(self, *, conninfo, **_kwargs):
|
||||
self.conninfo = conninfo
|
||||
self.closed = False
|
||||
self.borrowed = 0
|
||||
self.returned = 0
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
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
|
||||
@contextmanager
|
||||
def connection(self):
|
||||
self.borrowed += 1
|
||||
try:
|
||||
yield _FakeConnection()
|
||||
finally:
|
||||
self.returned += 1
|
||||
|
||||
|
||||
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 __init__(self):
|
||||
self.transactions = 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
|
||||
@contextmanager
|
||||
def transaction(self):
|
||||
self.transactions += 1
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_native_connections():
|
||||
connection.g_conn_dict.clear()
|
||||
connection.g_conninfo_dict.clear()
|
||||
connection._project_locks.clear()
|
||||
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.g_conn_dict.clear()
|
||||
connection.g_conninfo_dict.clear()
|
||||
connection._project_locks.clear()
|
||||
connection._pools.clear()
|
||||
connection._pool_conninfo.clear()
|
||||
connection._pool_borrows.clear()
|
||||
connection._admin_pools.clear()
|
||||
connection._admin_pool_borrows.clear()
|
||||
|
||||
|
||||
def test_is_project_open_drops_closed_cached_connection():
|
||||
connection.g_conn_dict["fengyang"] = _FakeConnection(closed=True)
|
||||
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}")
|
||||
|
||||
assert project.is_project_open("fengyang") is False
|
||||
assert "fengyang" not in connection.g_conn_dict
|
||||
first = connection.get_project_pool("fengyang")
|
||||
second = connection.get_project_pool("fengyang")
|
||||
|
||||
assert first is second
|
||||
assert first.conninfo == "dbname=fengyang"
|
||||
|
||||
|
||||
def test_open_connection_reuses_healthy_cached_connection(monkeypatch):
|
||||
cached = _FakeConnection()
|
||||
connection.g_conn_dict["fengyang"] = cached
|
||||
connection.g_conninfo_dict["fengyang"] = "dbname=fengyang"
|
||||
monkeypatch.setattr(
|
||||
connection, "get_project_pgconn_string", lambda db_name: f"dbname={db_name}"
|
||||
)
|
||||
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"])
|
||||
|
||||
def fail_connect(*, conninfo, autocommit):
|
||||
raise AssertionError("cached connection should be reused")
|
||||
old = connection.get_project_pool("fengyang")
|
||||
dsn["value"] = "host=new dbname=fengyang"
|
||||
new = connection.get_project_pool("fengyang")
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fail_connect)
|
||||
|
||||
assert connection.open_connection("fengyang") is cached
|
||||
assert cached.executed == ["SELECT 1"]
|
||||
assert old.closed is True
|
||||
assert new is not old
|
||||
assert new.conninfo == "host=new dbname=fengyang"
|
||||
|
||||
|
||||
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
|
||||
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}")
|
||||
|
||||
opened = []
|
||||
|
||||
def fake_connect(*, conninfo, autocommit):
|
||||
opened.append((conninfo, autocommit))
|
||||
return fresh
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fake_connect)
|
||||
monkeypatch.setattr(
|
||||
connection, "get_project_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"]
|
||||
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_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
|
||||
connection.g_conninfo_dict["fengyang"] = "dbname=fengyang"
|
||||
|
||||
opened = []
|
||||
|
||||
def fake_connect(*, conninfo, autocommit):
|
||||
opened.append((conninfo, autocommit))
|
||||
return fresh
|
||||
|
||||
monkeypatch.setattr(connection.pg, "connect", fake_connect)
|
||||
monkeypatch.setattr(
|
||||
connection, "get_project_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"]
|
||||
|
||||
|
||||
def test_open_connection_replaces_cache_when_project_dsn_changes(monkeypatch):
|
||||
cached = _FakeConnection()
|
||||
fresh = _FakeConnection()
|
||||
connection.g_conn_dict["fengyang"] = cached
|
||||
connection.g_conninfo_dict["fengyang"] = "host=old dbname=fengyang"
|
||||
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"host=new dbname={db_name}",
|
||||
lambda *, db_name: f"dbname={db_name}",
|
||||
)
|
||||
monkeypatch.setattr(connection.pg, "connect", lambda **_kwargs: fresh)
|
||||
|
||||
assert connection.open_connection("fengyang") is fresh
|
||||
assert cached.close_calls == 1
|
||||
assert connection.g_conninfo_dict["fengyang"] == "host=new dbname=fengyang"
|
||||
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
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
from app.native.wndb.core import database
|
||||
|
||||
|
||||
def _command(statement: str) -> database.DatabaseCommand:
|
||||
return database.DatabaseCommand(statement, [])
|
||||
|
||||
|
||||
def test_database_command_has_no_removed_undo_state() -> None:
|
||||
changes = [{"operation": "update", "type": "title", "value": "new"}]
|
||||
|
||||
command = database.DatabaseCommand("SELECT 1", changes)
|
||||
|
||||
assert vars(command) == {"sql": "SELECT 1", "changes": changes}
|
||||
assert not hasattr(command, "undo_sql")
|
||||
assert not hasattr(command, "undo_cs")
|
||||
|
||||
|
||||
def test_direct_model_write_refreshes_materialized_views(monkeypatch) -> None:
|
||||
events: list[str] = []
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"write",
|
||||
lambda _name, _statement: events.append("write"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"is_project_transaction_active",
|
||||
lambda _name: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"refresh_materialized_views",
|
||||
lambda _name: events.append("refresh"),
|
||||
)
|
||||
|
||||
result = database.execute_command(
|
||||
"project_a",
|
||||
database.DatabaseCommand(
|
||||
"UPDATE network.junctions SET elevation = 1",
|
||||
[{"operation": "update", "type": "junction", "id": "J-1"}],
|
||||
),
|
||||
)
|
||||
|
||||
assert events == ["write", "refresh"]
|
||||
assert result.operations == [
|
||||
{"operation": "update", "type": "junction", "id": "J-1"}
|
||||
]
|
||||
|
||||
|
||||
def test_batch_model_write_defers_materialized_view_refresh(monkeypatch) -> None:
|
||||
events: list[str] = []
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"write",
|
||||
lambda _name, _statement: events.append("write"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"is_project_transaction_active",
|
||||
lambda _name: True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"refresh_materialized_views",
|
||||
lambda _name: events.append("refresh"),
|
||||
)
|
||||
|
||||
database.execute_command(
|
||||
"project_a",
|
||||
_command("UPDATE network.junctions SET elevation = 1"),
|
||||
)
|
||||
|
||||
assert events == ["write"]
|
||||
|
||||
|
||||
def test_non_gis_model_write_does_not_refresh_materialized_views(monkeypatch) -> None:
|
||||
events: list[str] = []
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"write",
|
||||
lambda _name, _statement: events.append("write"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
database,
|
||||
"refresh_materialized_views",
|
||||
lambda _name: events.append("refresh"),
|
||||
)
|
||||
|
||||
database.execute_command(
|
||||
"project_a",
|
||||
_command("UPDATE network.time_settings SET value = '01:00'"),
|
||||
)
|
||||
|
||||
assert events == ["write"]
|
||||
@@ -1,4 +1,9 @@
|
||||
from app.native.wndb import s2_junctions
|
||||
import ast
|
||||
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.model import controls, junctions, patterns
|
||||
|
||||
|
||||
def test_get_junction_binds_untrusted_identifier(monkeypatch) -> None:
|
||||
@@ -9,13 +14,104 @@ def test_get_junction_binds_untrusted_identifier(monkeypatch) -> None:
|
||||
calls.append((name, statement, params))
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(s2_junctions, "try_read", fake_try_read)
|
||||
monkeypatch.setattr(junctions, "try_read", fake_try_read)
|
||||
|
||||
assert s2_junctions.get_junction("project_a", malicious_id) == {}
|
||||
assert calls == [
|
||||
(
|
||||
"project_a",
|
||||
"select * from junctions where id = %s",
|
||||
(malicious_id,),
|
||||
)
|
||||
]
|
||||
assert junctions.get_junction("project_a", malicious_id) == {}
|
||||
assert len(calls) == 1
|
||||
name, statement, params = calls[0]
|
||||
assert name == "project_a"
|
||||
assert "network.junctions" in statement
|
||||
assert "WHERE n.id = %s" in statement
|
||||
assert malicious_id not in statement
|
||||
assert params == (malicious_id,)
|
||||
|
||||
|
||||
def test_get_all_junctions_reads_materialized_view(monkeypatch) -> None:
|
||||
statements: list[str] = []
|
||||
|
||||
def fake_read_all(_name, statement):
|
||||
statements.append(statement)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(junctions, "read_all", fake_read_all)
|
||||
monkeypatch.setattr(junctions, "get_all_node_links", lambda _name: {})
|
||||
|
||||
assert junctions.get_all_junctions("project_a") == []
|
||||
assert "FROM gis.junctions" in statements[0]
|
||||
|
||||
|
||||
def test_get_all_scada_info_reads_materialized_view(monkeypatch) -> None:
|
||||
statements: list[str] = []
|
||||
|
||||
def fake_read_all(_name, statement):
|
||||
statements.append(statement)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(scada_assets, "read_all", fake_read_all)
|
||||
|
||||
assert scada_assets.get_all_scada_info("project_a") == []
|
||||
assert "FROM gis.scada_devices" in statements[0]
|
||||
|
||||
|
||||
def test_sql_literal_keeps_attacker_text_inside_one_postgres_literal() -> None:
|
||||
malicious = "x'); DROP SCHEMA network CASCADE; --"
|
||||
|
||||
assert sql_literal("O'Brien") == "'O''Brien'"
|
||||
assert sql_literal(malicious) == "'x''); DROP SCHEMA network CASCADE; --'"
|
||||
assert sql_literal(None) == "NULL"
|
||||
|
||||
|
||||
def test_pattern_command_quotes_malicious_identifier() -> None:
|
||||
malicious = "x'); DROP SCHEMA network CASCADE; --"
|
||||
change = ChangeSet({"id": malicious, "factors": [1.0]})
|
||||
|
||||
command = patterns._add_pattern("project_a", change).sql
|
||||
|
||||
assert f"values ({sql_literal(malicious)})" in command
|
||||
assert "values ('x'); DROP SCHEMA" not in command
|
||||
|
||||
|
||||
def test_inp_control_quotes_embedded_apostrophe() -> None:
|
||||
command = controls.inp_in_control("LINK P-1 STATUS 'OPEN'; DELETE")
|
||||
|
||||
assert "''OPEN''" in command
|
||||
assert "STATUS 'OPEN'; DELETE')" not in command
|
||||
|
||||
|
||||
def test_wndb_sql_fstrings_do_not_quote_formatted_values_directly() -> None:
|
||||
root = Path(__file__).resolve().parents[2] / "app" / "native" / "wndb"
|
||||
violations: list[str] = []
|
||||
for path in root.rglob("*.py"):
|
||||
tree = ast.parse(path.read_text(encoding="utf-8"))
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.JoinedStr):
|
||||
continue
|
||||
static_text = "".join(
|
||||
value.value
|
||||
for value in node.values
|
||||
if isinstance(value, ast.Constant) and isinstance(value.value, str)
|
||||
).lower()
|
||||
if not any(
|
||||
keyword in static_text
|
||||
for keyword in ("select ", "insert into", "update ", "delete from")
|
||||
):
|
||||
continue
|
||||
for index, value in enumerate(node.values):
|
||||
if not isinstance(value, ast.FormattedValue):
|
||||
continue
|
||||
before = node.values[index - 1] if index else None
|
||||
after = node.values[index + 1] if index + 1 < len(node.values) else None
|
||||
left_quote = (
|
||||
isinstance(before, ast.Constant)
|
||||
and isinstance(before.value, str)
|
||||
and before.value.endswith("'")
|
||||
)
|
||||
right_quote = (
|
||||
isinstance(after, ast.Constant)
|
||||
and isinstance(after.value, str)
|
||||
and after.value.startswith("'")
|
||||
)
|
||||
if left_quote and right_quote:
|
||||
violations.append(f"{path.name}:{node.lineno}")
|
||||
|
||||
assert violations == []
|
||||
|
||||
Reference in New Issue
Block a user