146 lines
5.9 KiB
Python
146 lines
5.9 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
|
|
|
|
|
|
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_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"
|