378 lines
19 KiB
Python
378 lines
19 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager, contextmanager
|
|
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 pydantic import ValidationError
|
|
|
|
from app.mcp.catalog import TOOLS
|
|
from app.mcp.models import BoundContext, QueryResult
|
|
from app.mcp import extended_service, mutation_service, query_service, scenario_service, analysis_service
|
|
|
|
|
|
HANDLE = "x" * 43
|
|
START = "2026-10-21T08:00:00+08:00"
|
|
|
|
|
|
def args(tool_name, **values):
|
|
return TOOLS[tool_name].model.model_validate({"context_handle": HANDLE, **values})
|
|
|
|
|
|
def binding(**values):
|
|
return BoundContext(context_handle=HANDLE, project_id=uuid4(), subject="actor", client_id="host",
|
|
native_session_id="session", purpose="ordinary_chat", tool_cap=list(TOOLS),
|
|
permission_cap=["webgis.edit", "webgis.view", "scada.clean", "scada.view", "simulation.run"],
|
|
expires_at=datetime.now(timezone.utc) + timedelta(minutes=5), **values)
|
|
|
|
|
|
@pytest.mark.parametrize("name,values", [
|
|
("query_network_extent", {"bounds": [120, 30, 122, 32]}),
|
|
("query_network_extent", {"bounds": [121, 31, 120, 30]}),
|
|
("query_network_extent", {"bounds": [float("nan"), 31, 122, 32]}),
|
|
("calculate_demand_allocation", {"target": "network", "demand": 10, "node_ids": ["J1"]}),
|
|
("edit_model_element", {"request_id": str(uuid4()), "operation": "delete", "element_type": "pipe", "element_id": "P1", "dry_run": False}),
|
|
("edit_model_element", {"request_id": str(uuid4()), "operation": "update", "element_type": "pipe", "element_id": "P1", "properties": {"id": "other"}}),
|
|
("clean_scada_readings", {"request_id": str(uuid4()), "start": START, "end": "2026-10-23T08:00:00+08:00", "device_ids": ["D1"]}),
|
|
("run_burst_simulation", {"start_time": START, "output_node_ids": ["J1"], "burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1, 2]}),
|
|
])
|
|
def test_bounded_schemas_reject_unsafe_inputs(name, values):
|
|
with pytest.raises(ValidationError):
|
|
args(name, **values)
|
|
|
|
|
|
def test_saved_scenario_has_no_nested_identity_and_validates_timezone():
|
|
values = dict(request_id=uuid4(), name="test", scenario="burst", parameters={
|
|
"start_time": START, "output_node_ids": ["J1"], "burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1],
|
|
})
|
|
valid = args("save_simulation_scenario", **values)
|
|
assert "context_handle" not in valid.parameters.model_dump()
|
|
values["parameters"]["start_time"] = "2026-10-21T08:00:00"
|
|
with pytest.raises(ValidationError):
|
|
args("save_simulation_scenario", **values)
|
|
values["parameters"]["start_time"] = START
|
|
values["parameters"]["network"] = "other-project"
|
|
with pytest.raises(ValidationError):
|
|
args("save_simulation_scenario", **values)
|
|
|
|
|
|
@pytest.mark.parametrize("purpose,role", [("task_reader", "member"), ("task_execution", "member"), ("task_collaboration", "member"), ("ordinary_chat", "viewer")])
|
|
def test_writes_never_enter_background_or_viewer_sessions(monkeypatch, purpose, role):
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"webgis.edit"})
|
|
bound = binding()
|
|
bound.purpose = purpose
|
|
scope = SimpleNamespace(project=SimpleNamespace(project_role=role))
|
|
with pytest.raises(HTTPException) as error:
|
|
asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))
|
|
assert error.value.status_code == 403
|
|
|
|
|
|
def test_write_requires_current_and_captured_permissions(monkeypatch):
|
|
scope = SimpleNamespace(project=SimpleNamespace(project_role="member"))
|
|
bound = binding()
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: set())
|
|
with pytest.raises(HTTPException):
|
|
asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))
|
|
monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"webgis.edit"})
|
|
bound.permission_cap = []
|
|
with pytest.raises(HTTPException):
|
|
asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))
|
|
|
|
|
|
def test_extent_runs_real_parameterized_project_query_and_pages(monkeypatch):
|
|
read = AsyncMock(return_value=[{"id": "A"}, {"id": "B"}])
|
|
monkeypatch.setattr(extended_service, "read_rows", read)
|
|
scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own"))
|
|
result = asyncio.run(extended_service.execute_extended(scope, object(), TOOLS["query_network_extent"],
|
|
args("query_network_extent", bounds=[121.0, 31.0, 121.1, 31.1], limit=1)))
|
|
assert result.page.has_more and result.truncated
|
|
assert result.data["items"] == [{"id": "A"}]
|
|
call = read.call_args.args
|
|
assert "ST_Intersects" in str(call[1])
|
|
assert call[2] == (121.0, 31.0, 121.1, 31.1, 2, 0)
|
|
assert call[0] is scope.connection
|
|
|
|
|
|
def test_missing_gis_propagates_unavailable_without_cross_project_retry(monkeypatch):
|
|
read = AsyncMock(side_effect=ValueError("GIS unavailable"))
|
|
monkeypatch.setattr(extended_service, "read_rows", read)
|
|
scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="no-gis"))
|
|
with pytest.raises(ValueError):
|
|
asyncio.run(extended_service.execute_extended(scope, object(), TOOLS["query_network_extent"],
|
|
args("query_network_extent", bounds=[121.0, 31.0, 121.1, 31.1])))
|
|
assert read.await_count == 1
|
|
|
|
|
|
def test_demand_allocation_is_paged_and_never_sets_demands(monkeypatch):
|
|
from app.services import demand_allocation
|
|
monkeypatch.setattr(demand_allocation, "calculate_demand_to_nodes", lambda *a: {"B": 6, "A": 4})
|
|
values = args("calculate_demand_allocation", target="nodes", node_ids=["A", "B"], demand=10, limit=1)
|
|
result = asyncio.run(analysis_service.calculate("calculate_demand_allocation", values, network="own", username="actor"))
|
|
assert result == {"items": [{"node_id": "A", "demand": 4}], "total": 2, "limit": 1, "offset": 0, "has_more": True, "persisted": False}
|
|
|
|
|
|
def test_burst_preview_uses_isolated_nonpersistent_adapter(monkeypatch):
|
|
from app.services import simulation_scenarios
|
|
calls = []
|
|
monkeypatch.setattr(simulation_scenarios, "burst_analysis", lambda **kw: calls.append(kw) or {"node_results": [{"node": "J1"}], "link_results": []})
|
|
values = args("run_burst_simulation", start_time=START, output_node_ids=["J1"], burst_pipe_ids=["P1"], orifice_areas_cm2=[2])
|
|
result = scenario_service.calculate_scenario("run_burst_simulation", values.model_dump(exclude={"context_handle", "timeout_seconds"}), network="own", username="actor")
|
|
assert calls[0]["persist_results"] is False and calls[0]["scheme_name"] is None
|
|
assert calls[0]["burst_size"] == [2]
|
|
assert result["persisted_baseline"] is False
|
|
|
|
|
|
class Connection:
|
|
@asynccontextmanager
|
|
async def transaction(self):
|
|
yield
|
|
|
|
async def commit(self):
|
|
pass
|
|
|
|
|
|
class Journal:
|
|
records = {}
|
|
|
|
def __init__(self, *args):
|
|
pass
|
|
|
|
async def lock_receipt(self, *args):
|
|
pass
|
|
|
|
async def receipt(self, scope, key):
|
|
return self.records.get((scope, key))
|
|
|
|
async def save_receipt(self, scope, key, parameters, result):
|
|
self.records[(scope, key)] = {"parameters": parameters, "result": result}
|
|
|
|
async def finish_receipt(self, scope, key, result):
|
|
self.records[(scope, key)]["result"] = result
|
|
|
|
|
|
@pytest.fixture
|
|
def mutation(monkeypatch):
|
|
Journal.records = {}
|
|
monkeypatch.setattr(mutation_service, "McpQueryRepository", Journal)
|
|
preview = AsyncMock(return_value={"before": {"diameter": 100}, "properties": {"diameter": 200}})
|
|
apply = AsyncMock(return_value={"changed": True})
|
|
monkeypatch.setattr(mutation_service, "preview_change", preview)
|
|
monkeypatch.setattr(mutation_service, "apply_change", apply)
|
|
scope = SimpleNamespace(connection=Connection(), project=SimpleNamespace(project_id=uuid4()))
|
|
values = dict(request_id=uuid4(), operation="update", element_type="pipe", element_id="P1", properties={"diameter": 200})
|
|
return scope, binding(), values, preview, apply
|
|
|
|
|
|
def test_preview_execution_stale_state_and_idempotent_replay(mutation):
|
|
scope, bound, values, preview, apply = mutation
|
|
tool = TOOLS["edit_model_element"]
|
|
async def run():
|
|
result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
|
|
assert result.data["dry_run"] and not apply.called
|
|
execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
|
|
preview.return_value = {"before": {"diameter": 150}}
|
|
with pytest.raises(HTTPException) as error:
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert error.value.status_code == 409 and not apply.called
|
|
assert Journal.records == {}
|
|
preview.return_value = {"before": {"diameter": 100}, "properties": {"diameter": 200}}
|
|
result = await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
replay = await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert replay == result
|
|
assert apply.await_count == 1
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_failed_or_timed_out_write_cannot_be_automatically_reexecuted(mutation):
|
|
scope, bound, values, preview, apply = mutation
|
|
tool = TOOLS["edit_model_element"]
|
|
async def run():
|
|
result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
|
|
execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
|
|
apply.side_effect = TimeoutError()
|
|
with pytest.raises(TimeoutError):
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
with pytest.raises(HTTPException) as error:
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert error.value.status_code == 409
|
|
assert apply.await_count == 1
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_capacity_rejection_is_retryable_under_the_same_request_after_slots_free(mutation, monkeypatch):
|
|
from threading import BoundedSemaphore
|
|
|
|
scope, bound, values, preview, apply = mutation
|
|
tool = TOOLS["edit_model_element"]
|
|
slots = BoundedSemaphore(2)
|
|
monkeypatch.setattr(analysis_service, "_slots", slots)
|
|
writes = []
|
|
|
|
async def bounded_write(*unused):
|
|
return await analysis_service.bounded_calculation(lambda: writes.append("written") or {"changed": True})
|
|
|
|
apply.side_effect = bounded_write
|
|
|
|
async def run():
|
|
result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
|
|
execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
|
|
assert slots.acquire(blocking=False) and slots.acquire(blocking=False)
|
|
try:
|
|
with pytest.raises(HTTPException) as error:
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert error.value.status_code == 503
|
|
record = next(iter(Journal.records.values()))
|
|
assert record["result"]["status"] == "rejected"
|
|
assert record["result"]["retryable"] is True
|
|
assert record["result"]["error"]["code"] == "CALCULATION_CAPACITY_EXCEEDED"
|
|
assert writes == []
|
|
finally:
|
|
slots.release()
|
|
slots.release()
|
|
saved = await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
replay = await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert replay == saved and writes == ["written"]
|
|
assert next(iter(Journal.records.values()))["result"]["status"] == "completed"
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.parametrize("failure", [TimeoutError(), asyncio.CancelledError(), RuntimeError("native failure"), HTTPException(503, "unrelated source failure")])
|
|
def test_ambiguous_write_failure_remains_started_and_cannot_be_retried(mutation, failure):
|
|
scope, bound, values, preview, apply = mutation
|
|
tool = TOOLS["edit_model_element"]
|
|
|
|
async def run():
|
|
result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
|
|
execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
|
|
apply.side_effect = failure
|
|
with pytest.raises(type(failure)):
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert next(iter(Journal.records.values()))["result"]["status"] == "started"
|
|
with pytest.raises(HTTPException) as error:
|
|
await mutation_service.execute_mutation(scope, bound, tool, execution)
|
|
assert error.value.status_code == 409 and apply.await_count == 1
|
|
asyncio.run(run())
|
|
|
|
|
|
def test_native_model_lock_revalidates_preview_and_only_writes_allowed_fields(monkeypatch):
|
|
from app.services import tjnetwork
|
|
from app.native.wndb.core import database
|
|
state = {"id": "P1", "diameter": 100}
|
|
monkeypatch.setattr(tjnetwork, "get_pipe", lambda *a: dict(state))
|
|
monkeypatch.setattr(tjnetwork, "get_pipe_schema", lambda *a: {"id": {"readonly": True}, "diameter": {"type": "float", "readonly": False}})
|
|
@contextmanager
|
|
def lock(*a):
|
|
yield
|
|
monkeypatch.setattr(database, "model_mutation_transaction", lock)
|
|
writes = []
|
|
monkeypatch.setattr(tjnetwork, "set_pipe", lambda network, changes: writes.append(changes.operations) or changes)
|
|
values = dict(request_id=uuid4(), operation="update", element_type="pipe", element_id="P1", properties={"diameter": 200})
|
|
preview = mutation_service.model_change("own", "edit_model_element", args("edit_model_element", **values), execute=False)
|
|
execution = args("edit_model_element", **values, dry_run=False, expected_digest=mutation_service.digest(preview))
|
|
state["diameter"] = 110
|
|
with pytest.raises(HTTPException):
|
|
mutation_service.model_change("own", "edit_model_element", execution, execute=True)
|
|
assert not writes
|
|
state["diameter"] = 100
|
|
mutation_service.model_change("own", "edit_model_element", execution, execute=True)
|
|
assert writes == [[{"id": "P1", "diameter": 200}]]
|
|
|
|
|
|
def test_artifacts_are_integrity_checked_bounded_binary_envelopes():
|
|
result = extended_service.artifact("model.zip", b"zip-data", "application/zip")
|
|
assert result["file"]["size"] == 8
|
|
assert len(result["file"]["sha256"]) == 64
|
|
with pytest.raises(HTTPException) as error:
|
|
extended_service.artifact("model.zip", bytes(5 * 1024 * 1024 + 1), "application/zip")
|
|
assert error.value.status_code == 413
|
|
|
|
|
|
def test_no_health_prediction_tool_and_every_write_requests_native_permission():
|
|
assert not any("health" in name for name in TOOLS)
|
|
assert len(TOOLS) == 67
|
|
assert sum(t.mutating for t in TOOLS.values()) == 13
|
|
for tool in TOOLS.values():
|
|
if tool.mutating:
|
|
assert not tool.protocol_tool().annotations.read_only_hint
|
|
assert tool.protocol_tool().annotations.destructive_hint
|
|
|
|
|
|
@pytest.mark.parametrize("scenario,parameters,function", [
|
|
("burst", {"burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1]}, "burst_analysis"),
|
|
("contaminant", {"source": "J2", "concentration": 1}, "contaminant_simulation"),
|
|
("valve_closure", {"valves": ["V1"]}, "valve_close_analysis"),
|
|
])
|
|
def test_formal_scenarios_bind_actor_and_persist_only_through_existing_services(monkeypatch, scenario, parameters, function):
|
|
from app.services import simulation_scenarios
|
|
calls = []
|
|
monkeypatch.setattr(simulation_scenarios, function, lambda **kw: calls.append(kw))
|
|
payload = {"start_time": START, "duration": 900, "output_node_ids": ["J1"], **parameters}
|
|
values = args("save_simulation_scenario", request_id=uuid4(), name="formal-test", scenario=scenario, parameters=payload)
|
|
mutation_service.saved_scenario("owned-project", "verified-user", values)
|
|
assert calls[0]["name"] == "owned-project"
|
|
assert calls[0]["username"] == "verified-user"
|
|
assert calls[0]["persist_results"] is True
|
|
assert calls[0]["scheme_name"] == "formal-test"
|
|
|
|
|
|
def test_mutation_status_is_scoped_to_bound_identity_session_and_project(monkeypatch):
|
|
calls = []
|
|
class Repository:
|
|
def __init__(self, conn, project):
|
|
calls.append(project)
|
|
async def receipt(self, scope_key, request_id):
|
|
calls.extend([scope_key, request_id])
|
|
return {"parameters": {"tool": "edit_model_element"}, "result": {"status": "started"}}
|
|
import app.infra.db.postgresql.mcp_queries as storage
|
|
monkeypatch.setattr(storage, "McpQueryRepository", Repository)
|
|
bound = binding()
|
|
request_id = uuid4()
|
|
scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own", project_id=bound.project_id))
|
|
result = asyncio.run(extended_service.execute_extended(scope, bound, TOOLS["get_mutation_status"], args("get_mutation_status", request_id=request_id)))
|
|
assert calls == [bound.project_id, "mutation:actor:session", str(request_id)]
|
|
assert result.data["status"] == "started"
|
|
|
|
|
|
def test_mutation_status_exposes_rejected_capacity_and_retryability(monkeypatch):
|
|
class Repository:
|
|
def __init__(self, *unused):
|
|
pass
|
|
async def receipt(self, *unused):
|
|
return {"parameters": {"tool": "save_simulation_scenario"}, "result": {
|
|
"status": "rejected", "retryable": True,
|
|
"error": {"code": "CALCULATION_CAPACITY_EXCEEDED", "message": "calculation not started"},
|
|
}}
|
|
import app.infra.db.postgresql.mcp_queries as storage
|
|
monkeypatch.setattr(storage, "McpQueryRepository", Repository)
|
|
bound = binding()
|
|
scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own", project_id=bound.project_id))
|
|
result = asyncio.run(extended_service.execute_extended(scope, bound, TOOLS["get_mutation_status"],
|
|
args("get_mutation_status", request_id=uuid4())))
|
|
assert result.data["status"] == "rejected" and result.data["retryable"] is True
|
|
assert result.data["error"]["code"] == "CALCULATION_CAPACITY_EXCEEDED"
|
|
|
|
|
|
def test_scada_preview_and_journal_distinguish_omitted_fields_from_explicit_null():
|
|
values = dict(request_id=uuid4(), operation="update", provenance="operator correction",
|
|
records=[{"device_id": "D1", "time": START, "monitored_value": 1}])
|
|
original = args("write_scada_readings", **values)
|
|
values["records"][0]["cleaned_value"] = None
|
|
clearing = args("write_scada_readings", **values)
|
|
for preview in [True, False]:
|
|
left = mutation_service.mutation_arguments(original, preview=preview)
|
|
right = mutation_service.mutation_arguments(clearing, preview=preview)
|
|
assert "cleaned_value" not in left["records"][0]
|
|
assert right["records"][0]["cleaned_value"] is None
|
|
assert mutation_service.digest(left) != mutation_service.digest(right)
|
|
|
|
|
|
def test_nested_model_properties_cannot_hide_nonfinite_or_unknown_fields():
|
|
schema = {"demands": {"type": "list", "readonly": False, "element": {"demand": {"type": "float", "readonly": False}}}}
|
|
for value in [{"demand": float("nan")}, {"demand": 1, "url": "https://foreign"}]:
|
|
with pytest.raises(HTTPException) as error:
|
|
mutation_service.check_properties(schema, {"demands": [value]})
|
|
assert error.value.status_code == 422
|