172 lines
8.3 KiB
Python
172 lines
8.3 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from app.mcp import query_service
|
|
from app.mcp.analysis_models import ScenarioInput
|
|
from app.mcp.catalog import TOOLS
|
|
from app.services import simulation
|
|
|
|
|
|
def parameters(**changes):
|
|
return ScenarioInput.model_validate({
|
|
"context_handle": "a" * 43, "start_time": "2026-07-16T08:00:00+08:00",
|
|
"duration": 60, "output_node_ids": ["J1"], **changes,
|
|
})
|
|
|
|
|
|
def context(monkeypatch):
|
|
permissions = {"simulation.run", "simulation.view", "scada.view", "webgis.view", "optimization.run"}
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions)
|
|
|
|
@asynccontextmanager
|
|
async def routing():
|
|
yield
|
|
|
|
scope = SimpleNamespace(
|
|
connection=SimpleNamespace(commit=AsyncMock()), simulation_routing=routing,
|
|
project=SimpleNamespace(project_code="owned", project_role="member"),
|
|
user=SimpleNamespace(username="actor"),
|
|
)
|
|
binding = SimpleNamespace(
|
|
tool_cap=["run_simulation"], permission_cap=sorted(permissions),
|
|
purpose="task_execution", task_id=uuid4(), task_scope_cap=[],
|
|
replay_context={"knowledge_cutoff": "2026-07-16T08:00:00+08:00"},
|
|
)
|
|
return scope, binding
|
|
|
|
|
|
def test_replay_coordinator_can_rebuild_online_baseline_without_live_scada(monkeypatch):
|
|
scope, binding = context(monkeypatch)
|
|
monkeypatch.setattr(simulation, "get_time", lambda _: {"HYDRAULIC TIMESTEP": "0:15:00"})
|
|
live_mappings = Mock(side_effect=AssertionError("Replay must not load live SCADA mappings"))
|
|
monkeypatch.setattr(simulation, "query_corresponding_element_id_and_query_id", live_mappings)
|
|
native = Mock(return_value={
|
|
"units": {"pressure": "MTR"},
|
|
"node_results": [{"node": "J1", "result": [{"pressure": 30}]}],
|
|
})
|
|
monkeypatch.setattr(simulation, "run_simulation", native)
|
|
|
|
result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_simulation"], parameters()))
|
|
|
|
assert result.origin == "synthetic"
|
|
assert any("不读取实时SCADA" in gap for gap in result.data_gaps)
|
|
assert result.data["result"]["persisted_baseline"] is True
|
|
assert result.data["result"]["node_results"][0]["result"] == [{"pressure": 30}]
|
|
native.assert_called_once()
|
|
call = native.call_args.kwargs
|
|
assert call["name"] == "owned" and call["simulation_type"] == "realtime"
|
|
assert call["modify_pattern_start_time"] == "2026-07-16T08:00:00+08:00"
|
|
assert all(not values for values in vars(call["scada_mappings"]).values())
|
|
live_mappings.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize("problem", ["future_start", "reader", "permission", "tool_cap"])
|
|
def test_replay_online_baseline_preserves_existing_authorization(monkeypatch, problem):
|
|
scope, binding = context(monkeypatch)
|
|
args = parameters()
|
|
if problem == "future_start":
|
|
args = parameters(start_time="2026-07-16T08:15:00+08:00")
|
|
if problem == "reader":
|
|
binding.purpose = "task_reader"
|
|
if problem == "permission":
|
|
binding.permission_cap.remove("simulation.run")
|
|
if problem == "tool_cap":
|
|
binding.tool_cap.clear()
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS["run_simulation"], args))
|
|
assert caught.value.status_code == (422 if problem == "future_start" else 403)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["analyze_valve_closure", "analyze_pressure_regulation"])
|
|
def test_replay_can_compare_hydraulic_alternatives_with_model_only_boundaries(monkeypatch, name):
|
|
from app.infra.db.postgresql.scada import ScadaElementMappings
|
|
from app.services import simulation_scenarios as scenarios
|
|
|
|
scope, binding = context(monkeypatch)
|
|
binding.tool_cap = [name]
|
|
captures = []
|
|
monkeypatch.setattr(scenarios, "get_option", lambda _: {"UNITS": "LPS"})
|
|
monkeypatch.setattr(scenarios, "set_option", lambda *_: None)
|
|
monkeypatch.setattr(simulation, "query_corresponding_element_id_and_query_id",
|
|
lambda _: ScadaElementMappings({}, {}, {"P1": "device"}, {}, {}, {}, {}))
|
|
monkeypatch.setattr(simulation, "run_simulation", lambda **kw: captures.append(kw) or {
|
|
"node_results": [{"node": "J1", "result": [{"pressure": 30}]}],
|
|
})
|
|
native_name = "valve_close_analysis" if name == "analyze_valve_closure" else "pressure_regulation"
|
|
native = getattr(scenarios, native_name).__wrapped__
|
|
monkeypatch.setattr(scenarios, native_name,
|
|
lambda **kw: native(**kw, _temporary_project="tjw_tmp_owned"))
|
|
args = TOOLS[name].model.model_validate({
|
|
**parameters().model_dump(),
|
|
**({"valves": ["V1"]} if name == "analyze_valve_closure" else {"pump_control": {"P1": [0]}}),
|
|
})
|
|
result = asyncio.run(query_service.execute_query(scope, binding, TOOLS[name], args))
|
|
assert result.origin == "synthetic"
|
|
assert result.data["result"]["persisted_baseline"] is False
|
|
assert len(captures) == 1
|
|
assert captures[0]["name"] == "tjw_tmp_owned" and captures[0]["persist_results"] is False
|
|
assert all(not values for values in vars(captures[0]["scada_mappings"]).values())
|
|
if name == "analyze_valve_closure":
|
|
assert captures[0]["modify_valve_opening"] == {"V1": 0.0}
|
|
else:
|
|
assert captures[0]["modify_fixed_pump_pattern"] == {"P1": [0]}
|
|
|
|
|
|
@pytest.mark.parametrize("name", [name for name, tool in TOOLS.items()
|
|
if tool.semantics == "asset" and not tool.mutating])
|
|
def test_replay_allows_authorized_asset_queries_for_input_discovery(monkeypatch, name):
|
|
scope, binding = context(monkeypatch)
|
|
binding.tool_cap = [name]
|
|
# Authorization needs only the schema kind; execution validates each tool's inputs.
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], SimpleNamespace(kind="node")))
|
|
binding.permission_cap = []
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], SimpleNamespace(kind="node")))
|
|
assert caught.value.status_code == 403
|
|
|
|
|
|
@pytest.mark.parametrize("name,fields", [
|
|
("query_model_baseline_series", {"element_type": "junction", "element_ids": ["J1"],
|
|
"start": "2026-07-16T07:45:00+08:00", "end": "2026-07-16T08:00:00+08:00"}),
|
|
("get_model_baseline_at_time", {"element_type": "junction", "element_id": "J1", "time": "2026-07-16T08:00:00+08:00"}),
|
|
("query_model_baseline_field", {"element_type": "junction", "field": "pressure", "time": "2026-07-16T08:00:00+08:00"}),
|
|
])
|
|
def test_replay_can_verify_baseline_without_reading_future_values(monkeypatch, name, fields):
|
|
scope, binding = context(monkeypatch)
|
|
binding.tool_cap = [name]
|
|
args = TOOLS[name].model.model_validate({"context_handle": "a" * 43, **fields})
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], args))
|
|
future = parameters(start_time="2026-07-16T08:15:00+08:00").start_time
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], args.model_copy(
|
|
update={"end" if "end" in fields else "time": future},
|
|
)))
|
|
assert caught.value.status_code == 422
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["calculate_demand_allocation", "optimize_sensor_placement"])
|
|
def test_replay_can_choose_current_model_design_calculations(monkeypatch, name):
|
|
scope, binding = context(monkeypatch)
|
|
binding.tool_cap = [name]
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], SimpleNamespace()))
|
|
binding.purpose = "task_reader"
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS[name], SimpleNamespace()))
|
|
assert caught.value.status_code == 403
|
|
|
|
|
|
def test_replay_can_explore_contamination_without_requiring_map_generation(monkeypatch):
|
|
scope, binding = context(monkeypatch)
|
|
binding.tool_cap = ["run_contaminant_simulation"]
|
|
args = TOOLS["run_contaminant_simulation"].model.model_validate({
|
|
**parameters().model_dump(), "source": "J1", "concentration": 1,
|
|
"include_impact_coverage": False,
|
|
})
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS["run_contaminant_simulation"], args))
|