200 lines
9.3 KiB
Python
200 lines
9.3 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from datetime import datetime, timedelta, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
from app.mcp.analysis_models import BurstLocationInput
|
|
from app.mcp.catalog import TOOLS
|
|
from app.mcp.models import BoundContext
|
|
from app.mcp import query_service
|
|
from app.services.scada_replay_queries import replay_burst_location_inputs, replay_snapshot_rows
|
|
|
|
|
|
START = datetime(2026, 10, 8, tzinfo=timezone.utc)
|
|
|
|
|
|
class Cursor:
|
|
def __init__(self, rows): self.rows = rows
|
|
async def fetchone(self): return self.rows[0] if self.rows else None
|
|
async def fetchall(self): return self.rows
|
|
|
|
|
|
class Connection:
|
|
def __init__(self, results):
|
|
self.results = list(results)
|
|
self.statements = []
|
|
|
|
@asynccontextmanager
|
|
async def transaction(self): yield
|
|
|
|
async def execute(self, statement, params):
|
|
self.statements.append((statement, params))
|
|
return Cursor(self.results.pop(0))
|
|
|
|
async def commit(self): pass
|
|
|
|
|
|
def _args(**changes):
|
|
return BurstLocationInput.model_validate({
|
|
"context_handle": "a" * 43,
|
|
"burst_leakage": 0.02,
|
|
"pressure_scada_ids": ["J1"],
|
|
"flow_scada_ids": ["L1"],
|
|
"scada_burst_start": START + timedelta(minutes=15),
|
|
"scada_burst_end": START + timedelta(minutes=30),
|
|
"scada_normal_start": START,
|
|
"scada_normal_end": START + timedelta(minutes=14),
|
|
**changes,
|
|
})
|
|
|
|
|
|
def test_replay_location_builds_series_only_from_frozen_outputs():
|
|
replay_id = str(uuid4())
|
|
run = {
|
|
"start_time": START.isoformat(),
|
|
"devices": [
|
|
{"device_id": "P1", "node_id": "J1", "link_id": None},
|
|
{"device_id": "F1", "node_id": None, "link_id": "L1"},
|
|
],
|
|
"monitoring_config": {
|
|
"pressure_device_ids": ["P1"],
|
|
"outlet_flow_device_ids": ["F1"],
|
|
},
|
|
}
|
|
readings = [
|
|
{"reading": {"device_id": "P1", "time": START.isoformat(), "cleaned_value": 100, "quality": "valid"}},
|
|
{"reading": {"device_id": "F1", "time": START.isoformat(), "cleaned_value": 50, "quality": "valid"}},
|
|
{"reading": {"device_id": "P1", "time": (START + timedelta(minutes=15)).isoformat(), "cleaned_value": 70, "quality": "valid"}},
|
|
{"reading": {"device_id": "F1", "time": (START + timedelta(minutes=15)).isoformat(), "cleaned_value": 80, "quality": "valid"}},
|
|
]
|
|
connection = Connection([[{"document": run}], readings])
|
|
context = {"replay_id": replay_id, "knowledge_cutoff": (START + timedelta(hours=1)).isoformat(),
|
|
"batch_sequence": 9}
|
|
|
|
result = asyncio.run(replay_burst_location_inputs(connection, uuid4(), context, _args()))
|
|
|
|
assert result["burst_pressure"] == {"J1": 70}
|
|
assert result["normal_pressure"] == {"J1": 100}
|
|
assert result["burst_flow"] == {"L1": 80}
|
|
assert result["normal_flow"] == {"L1": 50}
|
|
assert connection.statements[1][1][2] == 9
|
|
|
|
|
|
def test_replay_queries_frozen_history_with_project_replay_and_batch_bounds():
|
|
project_id, replay_id = uuid4(), str(uuid4())
|
|
at = START - timedelta(days=7)
|
|
run = {"start_time": START.isoformat(), "config": {"device_ids": ["P1"]},
|
|
"baseline": {"start_time": (START - timedelta(days=15)).isoformat()}}
|
|
reading = {"time": at.isoformat(), "device_id": "P1", "monitored_value": 23.,
|
|
"cleaned_value": 23., "quality": "valid", "reason": "冻结历史基线"}
|
|
connection = Connection([[{"document": run}], [{"reading": reading}]])
|
|
context = {"replay_id": replay_id, "knowledge_cutoff": START.isoformat(), "batch_sequence": 9}
|
|
result = asyncio.run(replay_snapshot_rows(connection, project_id, context, ["P1"], at, at, 10))
|
|
assert result == [reading]
|
|
statement, params = connection.statements[1]
|
|
assert "b.stage='detected'" in statement and "b.sequence<=%s" in statement
|
|
assert params[:4] == (str(project_id), replay_id, 9, ["P1"])
|
|
assert params[4:6] == (at, at)
|
|
|
|
|
|
@pytest.mark.parametrize("defect,status", [("old_replay", 422), ("too_early", 422), ("future", 422), ("foreign_device", 403)])
|
|
def test_replay_history_rejects_missing_freeze_or_scope_escape(defect, status):
|
|
at = START - timedelta(days=7)
|
|
run = {"start_time": START.isoformat(), "config": {"device_ids": ["P1"]},
|
|
"baseline": {"start_time": (START - timedelta(days=15)).isoformat()}}
|
|
if defect == "old_replay": run.pop("baseline")
|
|
if defect == "too_early": at = START - timedelta(days=16)
|
|
if defect == "future": at = START + timedelta(minutes=1)
|
|
connection = Connection([[{"document": run}]])
|
|
context = {"replay_id": str(uuid4()), "knowledge_cutoff": START.isoformat(), "batch_sequence": 9}
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(replay_snapshot_rows(connection, uuid4(), context,
|
|
["P2" if defect == "foreign_device" else "P1"], at, at, 10))
|
|
assert caught.value.status_code == status
|
|
assert len(connection.statements) <= 1
|
|
|
|
|
|
def test_location_accepts_frozen_pre_replay_normal_window():
|
|
replay_id = str(uuid4())
|
|
normal = START - timedelta(days=7)
|
|
run = {"start_time": START.isoformat(),
|
|
"baseline": {"start_time": (START - timedelta(days=15)).isoformat()},
|
|
"devices": [{"device_id": "P1", "node_id": "J1"}, {"device_id": "F1", "link_id": "L1"}],
|
|
"monitoring_config": {"pressure_device_ids": ["P1"], "outlet_flow_device_ids": ["F1"]}}
|
|
readings = [{"reading": {"time": at.isoformat(), "device_id": device,
|
|
"cleaned_value": value, "quality": "valid"}}
|
|
for at, device, value in [(normal, "P1", 100), (normal, "F1", 50),
|
|
(START + timedelta(minutes=15), "P1", 70),
|
|
(START + timedelta(minutes=15), "F1", 80)]]
|
|
connection = Connection([[{"document": run}], readings])
|
|
context = {"replay_id": replay_id, "knowledge_cutoff": (START + timedelta(hours=1)).isoformat(), "batch_sequence": 9}
|
|
result = asyncio.run(replay_burst_location_inputs(connection, uuid4(), context,
|
|
_args(scada_normal_start=normal, scada_normal_end=normal + timedelta(minutes=14))))
|
|
assert result["normal_pressure"] == {"J1": 100}
|
|
assert result["burst_flow"] == {"L1": 80}
|
|
|
|
|
|
def test_replay_location_authorization_requires_explicit_bounded_normal_window(monkeypatch):
|
|
binding = BoundContext(
|
|
context_handle="a" * 43, project_id=uuid4(), subject="u", client_id="host",
|
|
tool_cap=["run_burst_location"], permission_cap=["burst.run", "scada.view"],
|
|
native_session_id="s", purpose="task_execution", task_id=uuid4(),
|
|
expires_at=START + timedelta(days=1),
|
|
replay_context={"replay_id": str(uuid4()),
|
|
"knowledge_cutoff": (START + timedelta(hours=1)).isoformat(),
|
|
"batch_sequence": 9},
|
|
)
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"burst.run", "scada.view"})
|
|
scope = SimpleNamespace(project=SimpleNamespace(project_role="member"))
|
|
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS["run_burst_location"], _args()))
|
|
|
|
missing_normal = _args().model_copy(update={"scada_normal_start": None, "scada_normal_end": None})
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS["run_burst_location"], missing_normal))
|
|
assert caught.value.status_code == 422
|
|
|
|
|
|
def test_replay_location_execution_uses_frozen_adapter_instead_of_live_scada(monkeypatch):
|
|
from app.mcp import analysis_service
|
|
from app.services import scada_replay_queries
|
|
|
|
binding = BoundContext(
|
|
context_handle="a" * 43, project_id=uuid4(), subject="u", client_id="host",
|
|
tool_cap=["run_burst_location"], permission_cap=["burst.run", "scada.view"],
|
|
native_session_id="s", purpose="task_execution", task_id=uuid4(),
|
|
expires_at=START + timedelta(days=1),
|
|
replay_context={"replay_id": str(uuid4()),
|
|
"knowledge_cutoff": (START + timedelta(hours=1)).isoformat(),
|
|
"batch_sequence": 9},
|
|
)
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"burst.run", "scada.view"})
|
|
frozen = {"burst_leakage": 0.02, "burst_pressure": {"J1": 70}, "normal_pressure": {"J1": 100}}
|
|
build = AsyncMock(return_value=frozen)
|
|
calculate = AsyncMock(return_value={"candidates": [{"id": "PIPE-1"}]})
|
|
monkeypatch.setattr(scada_replay_queries, "replay_burst_location_inputs", build)
|
|
monkeypatch.setattr(analysis_service, "calculate_replay_burst_location", calculate)
|
|
|
|
@asynccontextmanager
|
|
async def routing(): yield
|
|
|
|
scope = SimpleNamespace(
|
|
connection=Connection([]),
|
|
project=SimpleNamespace(project_id=binding.project_id, project_code="demo", project_role="member"),
|
|
user=SimpleNamespace(username="tester"),
|
|
simulation_routing=routing,
|
|
)
|
|
result = asyncio.run(query_service.execute_query(
|
|
scope, binding, TOOLS["run_burst_location"], _args(),
|
|
))
|
|
|
|
build.assert_awaited_once()
|
|
calculate.assert_awaited_once_with(frozen, network="demo", username="tester")
|
|
assert result.origin == "synthetic"
|
|
assert result.data["result"]["candidates"][0]["id"] == "PIPE-1"
|