import asyncio from datetime import timedelta from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from app.mcp import calculation_jobs, query_service from app.mcp.analysis_models import BurstLocationInput from app.mcp.catalog import TOOLS from app.services import burst_location_demo as demo def arguments(**changes): return BurstLocationInput.model_validate({ "context_handle": "x" * 43, "burst_leakage": 0.02, "scada_burst_start": demo.EVENT_START - timedelta(minutes=15), "scada_burst_end": demo.EVENT_START, **changes, }) def topology(): def link(key, kind, a, b, length=0.0, diameter=None): return {"id": key, "kind": kind, "start_node_id": a, "end_node_id": b, "length": length, "diameter": diameter, "status": "OPEN"} return [ link(demo.REFERENCE_PIPE, "pipe", "170438", "J1", 350, 600), link(demo.CANDIDATE_PIPE, "pipe", "J1", "J2", 30, 600), link("V1", "valve", "S", "170438"), link("V2", "valve", "J2", "S"), {"id": "S", "kind": "reservoir", "start_node_id": None}, ] def test_candidate_is_on_same_line_within_distance_and_identical_valve_boundary(): result = demo._checked_candidate(topology()) assert result["pipe_id"] == demo.CANDIDATE_PIPE assert result["distance_m"] == 350 assert result["boundary_valve_ids"] == ["V1", "V2"] assert result["within_same_isolation_zone"] is True @pytest.mark.parametrize("defect", ["outside_zone", "different_line", "too_close", "too_far", "closed_line", "unisolatable", "missing_pipe"]) def test_unsuitable_candidate_is_never_returned(defect): rows = topology() if defect == "outside_zone": rows[1]["end_node_id"] = "S" elif defect == "different_line": rows[1]["diameter"] = 150 elif defect == "too_close": rows[0]["length"] = 199 elif defect == "too_far": rows[0]["length"] = 501 elif defect == "closed_line": rows[0]["status"] = "CLOSED" elif defect == "unisolatable": rows[2]["kind"] = "pipe" else: rows.pop(1) assert demo._checked_candidate(rows) is None @pytest.mark.parametrize("project,changes,enabled", [ ("tjwater2_v2", {}, True), ("tjw_bench_tjwater2_20261105", {}, True), ("tjwater_v2", {}, False), (None, {}, False), ("tjwater2_v2", {"scada_burst_start": demo.EVENT_START - timedelta(minutes=30), "scada_burst_end": demo.EVENT_START - timedelta(minutes=15)}, False), ("tjwater2_v2", {"scada_burst_end": demo.EVENT_END + timedelta(minutes=15)}, False), ]) def test_only_case3_demo_window_skips_native_calculation_queue(project, changes, enabled): args = arguments(**changes) assert demo.is_location_demo(project, args) is enabled assert calculation_jobs.should_queue(TOOLS["run_burst_location"], args, project_code=project) is not enabled def test_candidate_has_stable_similarity_without_presentation_blockers(monkeypatch): load = AsyncMock(return_value=topology()) monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", load) result = asyncio.run(demo.location_demo_result(object(), "tjwater2_v2", arguments())) assert result["located_pipe"] == demo.CANDIDATE_PIPE assert result["top_candidates"] == [{"pipe_id": demo.CANDIDATE_PIPE, "similarity": 0.9708, "similarity_percent": 97.08}] assert not {"mock", "notice", "simulation_times", "similarity_mode", "observed_source"} & result.keys() assert demo.REFERENCE_PIPE not in str(result) and demo.REFERENCE_NODE not in str(result) load.return_value = list(reversed(topology())) later = asyncio.run(demo.location_demo_result(object(), "tjwater2_v2", arguments( scada_burst_start=demo.EVENT_START, scada_burst_end=demo.EVENT_START + timedelta(minutes=15), ))) assert later == result load.return_value = [] empty = asyncio.run(demo.location_demo_result(object(), "tjwater2_v2", arguments())) assert empty["located_pipe"] is None and empty["top_candidates"] == [] assert empty["exit_condition"] == "no_candidate" assert "关阀范围" in empty["notice"] def test_mcp_demo_returns_synthetic_result_without_simulation_or_queue(monkeypatch): monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", AsyncMock(return_value=topology())) scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="tjwater2_v2")) binding = SimpleNamespace(replay_context=None) result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_burst_location"], arguments())) assert result.origin == "synthetic" assert result.data["result"]["located_pipe"] == demo.CANDIDATE_PIPE assert "mock" not in result.data["result"] assert result.data["result"]["top_candidates"][0]["similarity_percent"] == 97.08 assert result.data_gaps == [] def test_missing_candidate_preserves_a_specific_data_gap(monkeypatch): monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", AsyncMock(return_value=[])) scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="tjwater2_v2")) binding = SimpleNamespace(replay_context=None) result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_burst_location"], arguments())) assert result.data["result"]["top_candidates"] == [] assert result.data_gaps == [result.data["result"]["notice"]] assert "关阀范围" in result.data_gaps[0] def test_location_capability_and_tool_contract_do_not_inject_mock_instructions(): context = demo.location_demo_context("tjwater2_v2", {"knowledge_cutoff": demo.EVENT_START.isoformat()}) assert context["available"] is True assert "相似度" in context["notice"] and "影响分析" in context["notice"] assert "mock" not in str(context) and "模拟" not in context["notice"] tool = TOOLS["run_burst_location"].protocol_tool() assert "相似度" in tool.description assert "location_capability.available=true" in tool.description assert "location_demo" not in tool.description assert "demo模拟候选" not in tool.description assert "明确标记的演示候选" not in str(tool.input_schema) def test_replay_mock_still_validates_frozen_inputs_before_returning(monkeypatch): from fastapi import HTTPException from app.services import scada_replay_queries monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) freeze = AsyncMock(side_effect=HTTPException(422, "frozen baseline missing")) monkeypatch.setattr(scada_replay_queries, "replay_burst_location_inputs", freeze) load = AsyncMock() monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", load) scope = SimpleNamespace(connection=object(), project=SimpleNamespace( project_id="p", project_code="tjwater2_v2")) binding = SimpleNamespace(replay_context={"replay_id": "r"}) with pytest.raises(HTTPException): asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_burst_location"], arguments())) load.assert_not_awaited() def test_replay_mock_returns_after_frozen_validation_without_native_calculation(monkeypatch): from app.mcp import analysis_service from app.services import scada_replay_queries monkeypatch.setattr(query_service, "authorize_query", AsyncMock()) freeze = AsyncMock(return_value={"burst_pressure": {"P1": 20}, "normal_pressure": {"P1": 30}}) monkeypatch.setattr(scada_replay_queries, "replay_burst_location_inputs", freeze) native = AsyncMock(side_effect=AssertionError("mock entered native location")) monkeypatch.setattr(analysis_service, "calculate_replay_burst_location", native) monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", AsyncMock(return_value=topology())) scope = SimpleNamespace(connection=SimpleNamespace(commit=AsyncMock()), project=SimpleNamespace(project_id="p", project_code="tjwater2_v2")) binding = SimpleNamespace(replay_context={"replay_id": "r"}) result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_burst_location"], arguments())) assert result.origin == "synthetic" assert result.data["result"]["top_candidates"][0]["similarity"] == 0.9708 assert "mock" not in result.data["result"] assert result.data_gaps == [] freeze.assert_awaited_once() native.assert_not_awaited() def test_replay_demo_location_does_not_require_invented_leakage(monkeypatch): from app.mcp import analysis_service from app.services import scada_replay_queries permissions = {"burst.run", "scada.view"} monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions) args = arguments(burst_leakage=None, scada_normal_start=demo.EVENT_START - timedelta(days=1, minutes=15), scada_normal_end=demo.EVENT_START - timedelta(days=1)) context = {"replay_id": "r", "knowledge_cutoff": demo.EVENT_START.isoformat(), "batch_sequence": 3} binding = SimpleNamespace(replay_context=context, tool_cap=["run_burst_location"], permission_cap=permissions, purpose="task_execution", task_id="t") scope = SimpleNamespace(connection=SimpleNamespace(commit=AsyncMock()), project=SimpleNamespace(project_id="p", project_code="tjwater2_v2")) freeze = AsyncMock(return_value={"burst_pressure": {"J1": 20}, "normal_pressure": {"J1": 30}}) monkeypatch.setattr(scada_replay_queries, "replay_burst_location_inputs", freeze) monkeypatch.setattr(demo.NetworkAssetRepository, "get_isolation_model", AsyncMock(return_value=topology())) native = AsyncMock(side_effect=AssertionError("demo entered hydraulic calculation")) monkeypatch.setattr(analysis_service, "calculate_replay_burst_location", native) result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_burst_location"], args)) assert result.data["result"]["located_pipe"] == demo.CANDIDATE_PIPE assert result.data["result"]["burst_leakage"] is None assert result.origin == "synthetic" assert result.data["result"]["top_candidates"][0]["similarity_percent"] == 97.08 assert "mock" not in result.data["result"] freeze.assert_awaited_once_with(scope.connection, "p", context, args) native.assert_not_awaited() @pytest.mark.parametrize("project,changes", [ ("tjwater_v2", {}), ("tjwater2_v2", {"scada_burst_end": demo.EVENT_END + timedelta(minutes=15)}), ("tjwater2_v2", {"data_source": "simulation", "simulation_run_id": "8bfa5a5b-dbba-461a-88d4-34942bc66034"}), ]) def test_real_location_still_requires_leakage(monkeypatch, project, changes): from fastapi import HTTPException args = arguments(burst_leakage=None, **changes) permissions = {"burst.run", "scada.view", "simulation.view"} monkeypatch.setattr(query_service, "permissions_for_context", lambda _: permissions) binding = SimpleNamespace(replay_context=None, tool_cap=["run_burst_location"], permission_cap=permissions, purpose="ordinary_chat", task_id=None) with pytest.raises(HTTPException) as caught: asyncio.run(query_service.authorize_query(SimpleNamespace( project=SimpleNamespace(project_code=project)), binding, TOOLS["run_burst_location"], args)) assert caught.value.status_code == 422 assert "burst_leakage" in caught.value.detail