import asyncio from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import AsyncMock from uuid import uuid4 import pytest from fastapi import HTTPException from pydantic import ValidationError from app.domain.schemas.operational_data import TaskDataSnapshot from app.mcp import analysis_service, query_service from app.mcp.catalog import TOOLS from app.mcp.scope import QueryScope from app.infra.db.project_routing import ( ActiveProjectRouting, activate_project_routing, get_active_project_routing, get_project_pgconn_string, get_project_timescale_pgconn_string, ) HANDLE = "x" * 43 START = "2026-10-08T12:00:00+08:00" END = "2026-10-08T13:00:00+08:00" INPUTS = { "run_simulation": {"start_time": START, "duration": 3600, "output_node_ids": ["J1"], "output_link_ids": ["L1"]}, "run_contaminant_simulation": {"start_time": START, "output_node_ids": ["J1"], "source": "J2", "concentration": 1.0}, "analyze_valve_closure": {"start_time": START, "output_node_ids": ["J1"], "valves": ["V1"]}, "analyze_flushing": {"start_time": START, "output_node_ids": ["J1"], "drainage_node_id": "J2", "flushing_flow_m3h": 10}, "analyze_pressure_regulation": {"start_time": START, "output_node_ids": ["J1"], "pump_control": {"P1": [1]}}, "optimize_sensor_placement": {"method": "kmeans", "sensor_count": 3}, "run_burst_detection": {"target_time": START, "sensor_nodes": ["1", "2", "3", "4", "5"]}, "run_burst_location": {"data_source": "simulation", "simulation_run_id": str(uuid4()), "burst_leakage": 0.02, "pressure_scada_ids": ["1"], "scada_burst_start": START, "scada_burst_end": END}, "analyze_valve_isolation": {"accident_element": ["540672"]}, "run_leakage_identification": {"scada_start": START, "scada_end": END, "sensor_nodes": ["1"], "q_sum": 0.02, "start_time": 12, "dma_count": 1}, } def arguments(name, **extra): return TOOLS[name].model.model_validate({"context_handle": HANDLE, **INPUTS[name], **extra}) @pytest.fixture def routed_calculation_scope(monkeypatch): events = [] project = SimpleNamespace(project_id=uuid4(), project_code="logical-project") business = "host=localhost dbname=physical_business" timescale = "host=localhost dbname=physical_iot" async def resolve(project_id, role): assert project_id == project.project_id events.append(role) return SimpleNamespace( db_type="postgresql" if role == "biz_data" else "timescaledb", dsn=business if role == "biz_data" else timescale, ) async def commit(): events.append("metadata_commit") async def release_business_read(): events.append("business_commit") metadata = SimpleNamespace( get_project_db_routing=AsyncMock(side_effect=resolve), session=SimpleNamespace(commit=AsyncMock(side_effect=commit)), ) routing = ActiveProjectRouting(project.project_code, business, project_id=project.project_id) scope = QueryScope( project, SimpleNamespace(username="actor"), SimpleNamespace(commit=release_business_read), routing, metadata, ) monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) return scope, timescale, events @pytest.mark.parametrize("name", [ "run_simulation", "run_contaminant_simulation", "analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation", "run_burst_detection", "run_burst_location", "run_leakage_identification", ]) def test_timescale_calculation_routes_are_ready_before_native_worker(monkeypatch, routed_calculation_scope, name): scope, timescale, events = routed_calculation_scope def worker(tool_name, parameters, network, username): assert tool_name == name assert network == scope.project.project_code assert username == "actor" assert "context_handle" not in parameters assert "timeout_seconds" not in parameters assert get_project_timescale_pgconn_string(network) == timescale assert get_project_pgconn_string(network) == scope.routing.business_dsn assert events == ["business_commit", "biz_data", "iot_data", "metadata_commit"] events.append("worker") return {"routed": True} monkeypatch.setattr(analysis_service, "_calculate", worker) async def execute(): result = await query_service.execute_query(scope, object(), TOOLS[name], arguments(name)) assert get_active_project_routing() is scope.routing return result with activate_project_routing(scope.routing): result = asyncio.run(execute()) assert result.data["result"] == {"routed": True} assert events[-1] == "worker" def test_missing_timescale_route_stops_online_simulation_before_model_changes(monkeypatch, routed_calculation_scope): scope, _, _ = routed_calculation_scope resolve = scope.metadata.get_project_db_routing.side_effect async def missing_iot(project_id, role): return None if role == "iot_data" else await resolve(project_id, role) scope.metadata.get_project_db_routing.side_effect = missing_iot calculate = AsyncMock(return_value={}) monkeypatch.setattr(analysis_service, "calculate", calculate) with activate_project_routing(scope.routing): with pytest.raises(HTTPException) as denied: asyncio.run(query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation"))) assert get_active_project_routing() is scope.routing assert denied.value.status_code == 503 assert denied.value.detail == "Project TimescaleDB not configured" calculate.assert_not_awaited() @pytest.mark.parametrize("name", ["analyze_valve_isolation", "optimize_sensor_placement"]) def test_business_only_calculations_do_not_require_timescale(monkeypatch, routed_calculation_scope, name): scope, _, events = routed_calculation_scope def worker(*_): assert get_active_project_routing() is scope.routing assert get_project_pgconn_string(scope.project.project_code) == scope.routing.business_dsn return {"business_only": True} monkeypatch.setattr(analysis_service, "_calculate", worker) with activate_project_routing(scope.routing): result = asyncio.run(query_service.execute_query(scope, object(), TOOLS[name], arguments(name))) assert result.data["result"] == {"business_only": True} assert events == ["business_commit"] scope.metadata.get_project_db_routing.assert_not_awaited() def test_calculation_failure_restores_business_routing(monkeypatch, routed_calculation_scope): scope, timescale, _ = routed_calculation_scope def worker(*_): assert get_project_timescale_pgconn_string(scope.project.project_code) == timescale raise RuntimeError("native calculation failed") monkeypatch.setattr(analysis_service, "_calculate", worker) async def execute(): with pytest.raises(RuntimeError, match="native calculation failed"): await query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation")) assert get_active_project_routing() is scope.routing with activate_project_routing(scope.routing): asyncio.run(execute()) def test_denied_calculation_never_loads_timescale_or_starts_worker(monkeypatch, routed_calculation_scope): scope, _, events = routed_calculation_scope monkeypatch.setattr(query_service, "authorize_query", AsyncMock(side_effect=HTTPException(403, "denied"))) calculate = AsyncMock(return_value={}) monkeypatch.setattr(analysis_service, "calculate", calculate) with pytest.raises(HTTPException) as denied: asyncio.run(query_service.execute_query(scope, object(), TOOLS["run_simulation"], arguments("run_simulation"))) assert denied.value.status_code == 403 assert events == [] scope.metadata.get_project_db_routing.assert_not_awaited() calculate.assert_not_awaited() @pytest.mark.parametrize("name", INPUTS) def test_catalog_and_snapshots_accept_calculations_but_reject_model_owned_authority(name): tool = TOOLS[name] assert tool.computation assert arguments(name).timeout_seconds == 120 assert tool.protocol_tool().input_schema["additionalProperties"] is False for field in ("network", "username", "scheme_name", "final_candidates_csv_path", "observed_pressure_data", "persist_results", "simulation_type"): with pytest.raises(ValidationError): arguments(name, **{field: "untrusted"}) snapshot = TaskDataSnapshot(source=name, source_ref="mcp:test", queried_at=datetime.now(timezone.utc), origin="system_data", data={}, truncated=False) assert snapshot.source == name @pytest.mark.parametrize("name", INPUTS) def test_calculation_permission_and_reader_boundary(monkeypatch, name): permissions = {"burst.run", "simulation.run", "simulation.view", "scada.view", "optimization.run"} monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions) scope = SimpleNamespace(project=SimpleNamespace(project_role="member")) binding = SimpleNamespace(tool_cap=list(TOOLS), purpose="task_execution", permission_cap=list(permissions), task_id=uuid4(), task_scope_cap=["scada_readings"]) tool, args = TOOLS[name], arguments(name) asyncio.run(query_service.authorize_query(scope, binding, tool, args)) binding.purpose = "task_reader" with pytest.raises(HTTPException) as denied: asyncio.run(query_service.authorize_query(scope, binding, tool, args)) assert denied.value.status_code == 403 binding.purpose = "task_execution" binding.permission_cap.remove(tool.permission) with pytest.raises(HTTPException): asyncio.run(query_service.authorize_query(scope, binding, tool, args)) binding.permission_cap = list(permissions) permissions.remove(tool.permission) with pytest.raises(HTTPException): asyncio.run(query_service.authorize_query(scope, binding, tool, args)) def test_time_sources_and_finite_inputs_are_explicit(): invalid = [ ("run_simulation", {"start_time": "2026-10-08T12:00:00"}), ("run_simulation", {"duration": 0}), ("run_simulation", {"duration": 86401}), ("run_simulation", {"timeout_seconds": 0}), ("run_simulation", {"timeout_seconds": 86401}), ("run_simulation", {"timeout_seconds": 30.5}), ("run_simulation", {"timeout_seconds": True}), ("run_simulation", {"output_node_ids": []}), ("run_simulation", {"output_link_ids": [str(i) for i in range(21)]}), ("run_burst_detection", {"target_time": "2026-10-08T12:00:00"}), ("run_burst_detection", {"sensor_nodes": ["1"] * 5}), ("run_burst_location", {"burst_leakage": float("nan")}), ("run_burst_location", {"simulation_run_id": None}), ("run_burst_location", {"data_source": "monitoring"}), ("run_burst_location", {"scada_burst_end": "2026-10-10T12:00:00+08:00"}), ("run_leakage_identification", {"q_sum": 0}), ] for name, changes in invalid: with pytest.raises(ValidationError): arguments(name, **changes) args = arguments("run_burst_location", data_source="monitoring", simulation_run_id=None, scada_normal_start="2026-10-07T12:00:00+08:00", scada_normal_end="2026-10-07T13:00:00+08:00") assert args.data_source == "monitoring" @pytest.mark.parametrize("overrides", [{}, { "pressure_scada_ids": None, "scada_normal_start": None, "scada_normal_end": None, }, { "pressure_scada_ids": ["J2"], "scada_normal_start": "2026-10-06T12:00:00+08:00", "scada_normal_end": "2026-10-06T13:00:00+08:00", }]) @pytest.mark.parametrize("missing_normal", [False, True]) def test_burst_location_mcp_reuses_project_sensors_and_normal_window_defaults(monkeypatch, tmp_path, overrides, missing_normal): from app.services import burst_location args = TOOLS["run_burst_location"].model.model_validate({ "context_handle": HANDLE, "burst_leakage": 0.0274, "scada_burst_start": START, "scada_burst_end": END, **overrides, }) devices = [{"device_type": "pressure", "node_id": node, "api_query_id": f"q-{node}"} for node in ("J1", "J2")] monkeypatch.setattr(burst_location, "get_all_scada_info", lambda network: devices) queries = [] def read_scada(**kwargs): queries.append(kwargs) if missing_normal and len(queries) == 2: return {} return {device: [{"value": 20.0}] for device in kwargs["device_ids"]} monkeypatch.setattr(burst_location.InternalQueries, "query_scada_by_ids_timerange", read_scada) @contextmanager def inp(network, *, purpose): yield tmp_path / "model.inp" monkeypatch.setattr(burst_location, "temporary_project_inp", inp) algorithm = [] def locate(**kwargs): algorithm.append(kwargs) return {"located_pipe": "candidate"} monkeypatch.setattr(burst_location, "run_burst_location", locate) if missing_normal: with pytest.raises(ValueError, match="正常压力数据.*无有效数据"): asyncio.run(analysis_service.calculate( "run_burst_location", args, network="trusted-project", username="trusted-user", )) assert algorithm == [] return result = asyncio.run(analysis_service.calculate( "run_burst_location", args, network="trusted-project", username="trusted-user", )) expected_ids = overrides.get("pressure_scada_ids") or ["J1", "J2"] assert result["pressure_scada_ids"] == expected_ids assert [q["device_ids"] for q in queries] == [[f"q-{node}" for node in expected_ids]] * 2 normal_start = overrides.get("scada_normal_start") or "2026-10-07T12:00:00+08:00" normal_end = overrides.get("scada_normal_end") or "2026-10-07T13:00:00+08:00" assert datetime.fromisoformat(result["scada_window"]["normal_start"]) == datetime.fromisoformat(normal_start) assert datetime.fromisoformat(result["scada_window"]["normal_end"]) == datetime.fromisoformat(normal_end) assert algorithm[0]["pressure_scada_ids"] == expected_ids assert algorithm[0]["final_candidates_csv_path"] is None @pytest.mark.parametrize("overrides", [ {"scada_normal_start": "2026-10-07T12:00:00+08:00"}, {"scada_normal_end": "2026-10-07T13:00:00+08:00"}, {"scada_normal_start": "2026-10-07T12:00:00", "scada_normal_end": "2026-10-07T13:00:00"}, {"pressure_scada_ids": []}, {"pressure_scada_ids": ["J1", "J1"]}, {"pressure_scada_ids": [str(i) for i in range(21)]}, ]) def test_burst_location_optional_inputs_still_reject_partial_or_invalid_values(overrides): with pytest.raises(ValidationError): TOOLS["run_burst_location"].model.model_validate({ "context_handle": HANDLE, "burst_leakage": 0.0274, "scada_burst_start": START, "scada_burst_end": END, **overrides, }) @pytest.mark.parametrize("name", INPUTS) def test_calculations_accept_agent_execution_budget_separately_from_duration(name): args = arguments(name, timeout_seconds=900) assert args.timeout_seconds == 900 if hasattr(args, "duration"): assert args.duration == INPUTS[name].get("duration", 900) @pytest.mark.parametrize("name,module,function", [ ("run_burst_detection", "app.services.burst_detection", "run_burst_detection"), ("run_burst_location", "app.services.burst_location", "run_burst_location_by_network"), ("analyze_valve_isolation", "app.services.valve_isolation", "analyze_valve_isolation"), ("run_leakage_identification", "app.services.dma_leakage_estimation", "run_leakage_identification"), ]) def test_adapter_uses_existing_service_without_scheme_or_shared_output(monkeypatch, name, module, function): calls = [] def calculate(**kwargs): calls.append(kwargs) return {"ok": True} monkeypatch.setattr(f"{module}.{function}", calculate) result = asyncio.run(analysis_service.calculate(name, arguments(name), network="trusted-project", username="trusted-user")) assert result == {"ok": True} call = calls[0] assert call["network"] == "trusted-project" assert "context_handle" not in call if name != "analyze_valve_isolation": assert call["username"] == "trusted-user" assert call["scheme_name"] is None if name == "run_burst_location": assert call["final_candidates_csv_path"] is None if name == "run_leakage_identification": assert call["duration"] == 1 assert call["n_workers"] == 1 @pytest.mark.parametrize("name", ["run_burst_location", "run_simulation"]) def test_dispatch_ends_read_transaction_and_preserves_analysis_provenance(monkeypatch, name): order = [] async def commit(): order.append("commit") async def calculate(*args, **kwargs): order.append("calculate") assert kwargs == {"network": "trusted-project", "username": "trusted-user"} return {"located_pipe": "candidate"} @asynccontextmanager async def simulation_routing(): yield monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) monkeypatch.setattr(analysis_service, "calculate", calculate) scope = SimpleNamespace(connection=SimpleNamespace(commit=commit), project=SimpleNamespace(project_code="trusted-project"), user=SimpleNamespace(username="trusted-user"), simulation_routing=simulation_routing) result = asyncio.run(query_service.execute_query(scope, object(), TOOLS[name], arguments(name))) assert order == ["commit", "calculate"] assert result.semantics == "analysis" if name == "run_burst_location": assert result.data["parameters"]["burst_leakage"] == 0.02 else: assert result.data["parameters"]["duration"] == 3600 assert result.source == "run_simulation" assert "context_handle" not in result.data["parameters"] assert result.data["result"]["located_pipe"] == "candidate" assert "not observed SCADA" in result.data_gaps[-1] def test_busy_calculation_does_not_start_another_worker(monkeypatch): from threading import BoundedSemaphore slots = BoundedSemaphore(1) slots.acquire() monkeypatch.setattr(analysis_service, "_slots", slots) with pytest.raises(HTTPException) as denied: analysis_service._calculate("analyze_valve_isolation", {}, "project", "user") assert denied.value.status_code == 503 @pytest.mark.parametrize("name,function", [ ("run_contaminant_simulation", "contaminant_simulation"), ("analyze_valve_closure", "valve_close_analysis"), ("analyze_flushing", "flushing_analysis"), ("analyze_pressure_regulation", "pressure_regulation"), ]) def test_scenario_adapters_are_nonpersistent_and_return_only_requested_elements(monkeypatch, name, function): from app.mcp.scenario_service import calculate_scenario calls = [] def scenario(**kwargs): calls.append(kwargs) return {"node_results": [{"node": "J1", "result": []}, {"node": "J2", "result": []}], "link_results": [{"link": "L1", "result": []}]} monkeypatch.setattr(f"app.services.simulation_scenarios.{function}", scenario) monkeypatch.setattr("app.native.wndb.model.options.get_option", lambda _: {"UNITS": "CMH"}) monkeypatch.setattr("app.services.simulation.query_corresponding_element_id_and_query_id", lambda _: SimpleNamespace(fixed_pumps={"P1": "S1"}, variable_pumps={})) result = calculate_scenario(name, arguments(name).model_dump(exclude={"context_handle"}, exclude_none=True), network="trusted", username="actor") assert calls[0]["name"] == "trusted" assert calls[0]["persist_results"] is False assert calls[0]["scheme_name"] is None assert result["node_results"] == [{"node": "J1", "result": []}] assert result["link_results"] == [] assert result["coverage"]["missing_nodes"] == [] if name == "analyze_flushing": assert calls[0]["flushing_flow"] == 10 def test_sensor_preview_validates_but_does_not_store(monkeypatch): from app.services import sensor_placement as service from unittest.mock import Mock validate = Mock() save = Mock() monkeypatch.setattr(service, "validate_sensor_placement_nodes", validate) monkeypatch.setattr(service, "_sensor_points", lambda *_: [{"node_id": "J1"}]) monkeypatch.setattr(service.sensor_placement_repository, "create_sensor_placement", save) result = service._create_validated_placement("project", run_name="preview", min_diameter=0, created_by="actor", sensor_locations=["J1"], persist_results=False) validate.assert_called_once_with("project", ["J1"]) save.assert_not_called() assert result["sensor_locations"] == ["J1"] @pytest.mark.parametrize("name", ["run_contaminant_simulation", "run_simulation"]) def test_scenario_also_requires_scada_permission_even_with_task_read_scope(monkeypatch, name): permissions = {"simulation.run"} monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions) binding = SimpleNamespace(tool_cap=list(TOOLS), purpose="task_execution", permission_cap=list(permissions), task_id=uuid4(), task_scope_cap=["scada_readings"]) with pytest.raises(HTTPException): asyncio.run(query_service.authorize_query(SimpleNamespace(project=object()), binding, TOOLS[name], arguments(name))) @pytest.mark.parametrize("fails", [False, True]) def test_standard_simulation_uses_online_service_and_propagates_failure(monkeypatch, fails): from app.mcp.scenario_service import calculate_scenario events = [] def temporary_project(*_): pytest.fail("Online calculation must not copy a database") def simulate(network, start, duration, **kwargs): assert network == "trusted-project" assert start.isoformat() == START assert duration == 3600 assert kwargs == {"output_node_ids": ["J1"], "output_link_ids": ["L1"]} events.append("calculate") if fails: raise RuntimeError("calculation failed") return {"node_results": [{"node": "J1"}, {"node": "J2"}], "link_results": [{"link": "L1"}, {"link": "L2"}], "times": {"report_step": 900}, "units": {"pressure": "m", "flow": "LPS"}} monkeypatch.setattr("app.native.wndb.core.projects.temporary_project_database", temporary_project) monkeypatch.setattr("app.services.simulation.run_online_simulation", simulate) def execute(): return calculate_scenario("run_simulation", arguments("run_simulation").model_dump( exclude={"context_handle"}), network="trusted-project", username="actor") if fails: with pytest.raises(RuntimeError, match="calculation failed"): execute() else: result = execute() assert result["node_results"] == [{"node": "J1"}] assert result["link_results"] == [{"link": "L1"}] assert result["duration_seconds"] == 3600 assert result["persisted_scheme"] is False assert result["persisted_baseline"] is True assert result["units"]["pressure"] == "m" assert events == ["calculate"]