71 lines
3.4 KiB
Python
71 lines
3.4 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
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.mcp import analysis_service, query_service
|
|
from app.mcp.analysis_models import ContaminantSimulationInput
|
|
from app.mcp.calculation_jobs import should_queue
|
|
from app.mcp.catalog import TOOLS
|
|
|
|
|
|
def args(**extra):
|
|
return ContaminantSimulationInput(context_handle="a" * 43, **({
|
|
"start_time": "2026-07-16T08:00:00+08:00", "duration": 1800,
|
|
"output_node_ids": ["Q1"], "source": "Q1", "concentration": 1,
|
|
"include_impact_coverage": True, "impact_threshold_mg_l": 0.01,
|
|
} | extra))
|
|
|
|
|
|
@pytest.mark.parametrize("extra", [{"impact_threshold_mg_l": -1}, {"impact_threshold_mg_l": float("nan")},
|
|
{"include_impact_coverage": "true"}])
|
|
def test_impact_inputs_reject_invalid_thresholds_and_flags(extra):
|
|
with pytest.raises(ValidationError): args(**extra)
|
|
|
|
|
|
def context(monkeypatch):
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"simulation.run", "scada.view"})
|
|
@asynccontextmanager
|
|
async def routing():
|
|
yield
|
|
scope = SimpleNamespace(connection=SimpleNamespace(commit=AsyncMock()),
|
|
project=SimpleNamespace(project_code="project", project_role="member"), user=SimpleNamespace(username="actor"),
|
|
simulation_routing=routing)
|
|
binding = SimpleNamespace(tool_cap=["run_contaminant_simulation"], permission_cap=["simulation.run", "scada.view"],
|
|
purpose="task_execution", task_id=uuid4(), task_scope_cap=[],
|
|
replay_context={"knowledge_cutoff": "2026-07-16T08:00:00+08:00"})
|
|
return scope, binding
|
|
|
|
|
|
@pytest.mark.parametrize("include_impact", [True, False])
|
|
def test_replay_prediction_can_extend_past_cutoff_without_reading_future_observations(monkeypatch, include_impact):
|
|
scope, binding = context(monkeypatch)
|
|
worker = AsyncMock(return_value={"impact": {"affected_node_count": 2}})
|
|
monkeypatch.setattr(analysis_service, "calculate_replay_contaminant_simulation", worker)
|
|
parameters = args(include_impact_coverage=include_impact)
|
|
result = asyncio.run(query_service.execute_query(scope, binding, TOOLS["run_contaminant_simulation"], parameters))
|
|
assert result.origin == "synthetic" and result.data["result"]["impact"]["affected_node_count"] == 2
|
|
assert "预测" in result.data_gaps[0]
|
|
scope.connection.commit.assert_awaited_once()
|
|
worker.assert_awaited_once()
|
|
assert worker.await_args.args[0].include_impact_coverage is include_impact
|
|
assert worker.await_args.kwargs == {"network": "project", "username": "actor"}
|
|
assert should_queue(TOOLS["run_contaminant_simulation"], args())
|
|
|
|
|
|
@pytest.mark.parametrize("problem", ["future_start", "reader", "permission"])
|
|
def test_replay_admission_keeps_time_permission_and_reader_boundaries(monkeypatch, problem):
|
|
scope, binding = context(monkeypatch)
|
|
parameters = args()
|
|
if problem == "future_start": parameters = args(start_time="2026-07-16T08:15:00+08:00")
|
|
if problem == "reader": binding.purpose = "task_reader"
|
|
if problem == "permission": binding.permission_cap.remove("scada.view")
|
|
with pytest.raises(HTTPException) as caught:
|
|
asyncio.run(query_service.authorize_query(scope, binding, TOOLS["run_contaminant_simulation"], parameters))
|
|
assert caught.value.status_code == (422 if problem == "future_start" else 403)
|