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:
@@ -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