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))