from __future__ import annotations
import asyncio
import json
import threading
import time
import unittest
from unittest.mock import patch
from uuid import uuid4
from fastapi import HTTPException, Request
from app.main import (
ReactFlowEdgePayload,
ReactFlowProjectPayload,
SimulationCancellationPayload,
_mark_simulation_task_result,
_register_simulation_task,
build_reactflow_system_xml,
cancel_system_xml_simulation,
compile_reactflow_network,
get_system_xml_simulation,
simulation_event_stream,
simulate_system_xml,
)
from PythonModels.components.resistive_pipe import ResistivePipe
from PythonModels.core.solver import SolveIVPConfig
from PythonModels.systems.generic import (
GenericFluidSystem,
SimulationPreparationError,
)
from tests.test_system_xml_protocol import physical_port
def component_node(
component_id: str,
model_type: str,
ports: list[dict[str, str]],
parameters: dict[str, float] | None = None,
) -> dict[str, object]:
return {
"id": component_id,
"type": "simulationComponent",
"position": {"x": 0.0, "y": 0.0},
"data": {
"label": component_id,
"componentType": model_type,
"modelType": model_type,
"ports": ports,
"parameters": parameters or {},
},
}
def physical_edge(
edge_id: str,
first_component: str,
first_port: str,
second_component: str,
second_port: str,
) -> dict[str, str]:
return {
"id": edge_id,
"source": first_component,
"sourceHandle": first_port,
"target": second_component,
"targetHandle": second_port,
}
def two_port_definitions() -> list[dict[str, str]]:
return [
physical_port("port_a", "inlet", "left"),
physical_port("port_b", "outlet", "right"),
]
def tee_port_definitions() -> list[dict[str, str]]:
return [
physical_port("port_in", "bidirectional", "left"),
physical_port("port_out1", "bidirectional", "right"),
physical_port("port_out2", "bidirectional", "right"),
]
def chain_project(*, reverse_edges: bool = False) -> ReactFlowProjectPayload:
edges = [
physical_edge("edge-1", "cylinder_1", "port_b", "orifice_1", "port_a"),
physical_edge("edge-2", "orifice_1", "port_b", "pipe_1", "port_a"),
physical_edge("edge-3", "pipe_1", "port_b", "tank_1", "port_a"),
]
if reverse_edges:
edges = [
physical_edge(
edge["id"],
edge["target"],
edge["targetHandle"],
edge["source"],
edge["sourceHandle"],
)
for edge in edges
]
return ReactFlowProjectPayload(
name="generic-chain",
nodes=[
component_node(
"cylinder_1",
"cylinder",
[physical_port("port_b", "outlet", "right")],
{"volume": 0.01, "p0": 500000.0, "T0": 300.0},
),
component_node(
"orifice_1",
"orifice",
two_port_definitions(),
{"K": 1e-5, "opening": 1.0},
),
component_node(
"pipe_1",
"pipe",
two_port_definitions(),
{
"length": 1.0,
"diameter": 0.02,
"lambda_darcy": 0.02,
"p0": 100000.0,
"T0": 300.0,
},
),
component_node(
"tank_1",
"tank",
[physical_port("port_a", "inlet", "left")],
{"volume": 0.1, "p0": 100000.0, "T0": 300.0},
),
],
edges=edges,
simulation={
"t_start": 0.0,
"t_stop": 0.01,
"step": 0.005,
"max_step": 0.001,
"method": "BDF",
},
)
def branched_project() -> ReactFlowProjectPayload:
return ReactFlowProjectPayload(
name="generic-branch",
nodes=[
component_node(
"cylinder_1",
"cylinder",
[physical_port("port_b", "outlet", "right")],
{"volume": 0.01, "p0": 500000.0, "T0": 300.0},
),
component_node("tee_1", "tee", tee_port_definitions()),
component_node(
"orifice_1",
"orifice",
two_port_definitions(),
{"K": 1e-5, "opening": 1.0},
),
component_node(
"pipe_1",
"pipe",
two_port_definitions(),
{
"length": 1.0,
"diameter": 0.02,
"lambda_darcy": 0.02,
"p0": 100000.0,
"T0": 300.0,
},
),
component_node(
"orifice_2",
"orifice",
two_port_definitions(),
{"K": 1e-5, "opening": 1.0},
),
component_node(
"pipe_2",
"pipe",
two_port_definitions(),
{
"length": 1.0,
"diameter": 0.02,
"lambda_darcy": 0.02,
"p0": 100000.0,
"T0": 300.0,
},
),
component_node("tee_2", "tee", tee_port_definitions()),
component_node(
"tank_1",
"tank",
[physical_port("port_a", "inlet", "left")],
{"volume": 0.1, "p0": 100000.0, "T0": 300.0},
),
],
edges=[
physical_edge("edge-1", "cylinder_1", "port_b", "tee_1", "port_in"),
physical_edge("edge-2", "tee_1", "port_out1", "orifice_1", "port_a"),
physical_edge("edge-3", "orifice_1", "port_b", "pipe_1", "port_a"),
physical_edge("edge-4", "pipe_1", "port_b", "tee_2", "port_out1"),
physical_edge("edge-5", "tee_1", "port_out2", "orifice_2", "port_a"),
physical_edge("edge-6", "orifice_2", "port_b", "pipe_2", "port_a"),
physical_edge("edge-7", "pipe_2", "port_b", "tee_2", "port_out2"),
physical_edge("edge-8", "tee_2", "port_in", "tank_1", "port_a"),
],
simulation={
"t_start": 0.0,
"t_stop": 0.005,
"step": 0.005,
"max_step": 0.001,
"method": "BDF",
},
)
def xml_request(body: bytes) -> Request:
delivered = False
async def receive():
nonlocal delivered
if delivered:
return {"type": "http.disconnect"}
delivered = True
return {"type": "http.request", "body": body, "more_body": False}
return Request(
{
"type": "http",
"method": "POST",
"path": "/api/system-xml/simulate",
"headers": [(b"content-type", b"application/xml")],
},
receive,
)
class GenericSystemXmlSimulationTests(unittest.TestCase):
def test_xml_pipe_compiles_to_quasi_steady_resistance(self) -> None:
network = compile_reactflow_network(chain_project())
self.assertIsInstance(network.components["pipe_1"], ResistivePipe)
structure = network.pressure_flow_structure_dict()
self.assertEqual(structure["unknownCount"], 12)
self.assertEqual(structure["equationCount"], 12)
self.assertTrue(structure["isSquare"])
def test_generic_chain_simulation_conserves_mass_and_moves_pressures(self) -> None:
network = compile_reactflow_network(chain_project())
progress: list[tuple[float, str]] = []
result = GenericFluidSystem(network).simulate(
SolveIVPConfig(
t_start=0.0,
t_stop=0.01,
method="BDF",
max_step=0.001,
),
sample_step=0.005,
progress_callback=lambda value, phase: progress.append((value, phase)),
)
self.assertTrue(result.success)
self.assertEqual(progress[0], (0.0, "initializing"))
self.assertEqual(progress[-1], (1.0, "complete"))
self.assertTrue(
all(
current[0] <= following[0]
for current, following in zip(progress, progress[1:])
)
)
self.assertIn("integrating", {phase for _, phase in progress})
self.assertIn("postprocessing", {phase for _, phase in progress})
self.assertLess(result.series["cylinder_1.p"][-1], 500000.0)
self.assertGreater(result.series["tank_1.p"][-1], 100000.0)
total_mass = [
cylinder + tank
for cylinder, tank in zip(
result.series["cylinder_1.m"],
result.series["tank_1.m"],
)
]
self.assertLess(max(total_mass) - min(total_mass), 1e-12)
total_energy = [
cylinder + tank
for cylinder, tank in zip(
result.series["cylinder_1.U"],
result.series["tank_1.U"],
)
]
self.assertLess(max(total_energy) - min(total_energy), 1e-6)
self.assertLess(
result.diagnostics["pressureFlow"]["maxScaledResidual"],
1e-7,
)
def test_cancelled_simulation_returns_accepted_partial_samples(self) -> None:
cancel_event = threading.Event()
def request_cancel_after_progress(progress: float, phase: str) -> None:
if phase == "integrating" and progress >= 0.2:
cancel_event.set()
result = GenericFluidSystem(
compile_reactflow_network(chain_project())
).simulate(
SolveIVPConfig(t_stop=0.05, method="BDF", max_step=0.001),
sample_step=0.005,
progress_callback=request_cancel_after_progress,
cancel_check=cancel_event.is_set,
)
self.assertFalse(result.success)
self.assertEqual(result.status, "cancelled")
self.assertGreaterEqual(result.diagnostics["sampleCount"], 2)
self.assertGreater(result.simulated_until, 0.0)
self.assertLess(result.simulated_until, 0.05)
self.assertEqual(
len(result.series["time"]),
len(result.series["cylinder_1.p"]),
)
def test_task_registry_distinguishes_user_stop_and_stalled_stop(self) -> None:
for reason, expected_status in (("user", "stopped"), ("stalled", "stalled")):
simulation_id = f"test-{uuid4().hex}"
task = _register_simulation_task(simulation_id)
response = cancel_system_xml_simulation(
simulation_id,
SimulationCancellationPayload(reason=reason),
)
self.assertTrue(response["accepted"])
result = _mark_simulation_task_result(
task,
{
"success": False,
"status": "cancelled",
"partial": True,
"series": {"time": [0.0, 0.1]},
},
)
snapshot = get_system_xml_simulation(simulation_id)
self.assertEqual(result["status"], expected_status)
self.assertEqual(snapshot["status"], expected_status)
self.assertEqual(snapshot["cancelReason"], reason)
self.assertEqual(snapshot["result"]["series"]["time"], [0.0, 0.1])
def test_physical_edge_order_does_not_change_simulation(self) -> None:
forward = GenericFluidSystem(
compile_reactflow_network(chain_project())
).simulate(
SolveIVPConfig(t_stop=0.005, method="BDF", max_step=0.001),
sample_step=0.005,
)
reverse = GenericFluidSystem(
compile_reactflow_network(chain_project(reverse_edges=True))
).simulate(
SolveIVPConfig(t_stop=0.005, method="BDF", max_step=0.001),
sample_step=0.005,
)
self.assertAlmostEqual(
forward.final["tank_1.p"],
reverse.final["tank_1.p"],
places=7,
)
def test_branched_topology_is_solved_without_fixed_testmodel_closure(self) -> None:
network = compile_reactflow_network(branched_project())
result = GenericFluidSystem(network).simulate(
SolveIVPConfig(t_stop=0.005, method="BDF", max_step=0.001),
sample_step=0.005,
)
self.assertTrue(result.success)
self.assertAlmostEqual(
result.final["pipe_1.port_a.m_flow"],
result.final["pipe_2.port_a.m_flow"],
places=10,
)
self.assertGreater(result.final["tank_1.p"], 100000.0)
def test_directly_coupled_storage_components_are_rejected(self) -> None:
project = chain_project()
project.nodes = [project.nodes[0], project.nodes[-1]]
project.edges = [
ReactFlowEdgePayload(
**physical_edge(
"edge-1",
"cylinder_1",
"port_b",
"tank_1",
"port_a",
)
)
]
with self.assertRaises(SimulationPreparationError) as caught:
GenericFluidSystem(compile_reactflow_network(project))
self.assertIn(
"IDEAL_STORAGE_COUPLING_UNSUPPORTED",
{issue.code for issue in caught.exception.issues},
)
def test_raw_system_xml_runs_through_generic_simulation_endpoint(self) -> None:
xml = build_reactflow_system_xml(chain_project())
response = asyncio.run(simulate_system_xml(xml_request(xml)))
self.assertTrue(response["success"])
self.assertEqual(response["simulation"]["method"], "BDF")
self.assertEqual(response["model"]["pressureFlowSystem"]["unknownCount"], 12)
self.assertEqual(response["diagnostics"]["stateCount"], 4)
self.assertGreater(response["final"]["tank_1.p"], 100000.0)
variables = {variable["key"]: variable for variable in response["variables"]}
self.assertEqual(set(variables), set(response["series"]) - {"time"})
self.assertEqual(variables["cylinder_1.p"]["componentId"], "cylinder_1")
self.assertEqual(variables["cylinder_1.p"]["label"], "压力")
self.assertEqual(variables["cylinder_1.p"]["unit"], "Pa")
self.assertEqual(variables["pipe_1.port_a.m_flow"]["scope"], "port")
self.assertEqual(
variables["pipe_1.port_a.m_flow"]["portName"],
"port_a",
)
def test_streaming_endpoint_events_have_monotonic_progress_and_result(self) -> None:
project = chain_project()
xml = build_reactflow_system_xml(project)
events = [json.loads(line) for line in simulation_event_stream(xml)]
progress_events = [
event for event in events if event["event"] == "progress"
]
progress = [event["progress"] for event in progress_events]
self.assertGreater(len(progress), 3)
self.assertTrue(
all(current <= following for current, following in zip(progress, progress[1:]))
)
preparation_events = [
event
for event in progress_events
if event["phase"] in {"validation", "compilation", "initialization"}
]
self.assertTrue(preparation_events)
self.assertTrue(all(event["progress"] == 0 for event in preparation_events))
timed_events = [
event
for event in progress_events
if event.get("simulatedTime") is not None
]
self.assertGreater(len(timed_events), 1)
for event in timed_events:
self.assertEqual(event["totalTime"], project.simulation.t_stop)
expected_progress = round(
100
* (event["simulatedTime"] - project.simulation.t_start)
/ (project.simulation.t_stop - project.simulation.t_start)
)
self.assertEqual(event["progress"], expected_progress)
self.assertEqual(events[-1]["event"], "result")
self.assertEqual(events[-1]["progress"], 100)
self.assertTrue(events[-1]["result"]["success"])
def test_streaming_task_stop_returns_partial_result(self) -> None:
project = branched_project()
project.simulation.t_stop = 2.0
project.simulation.step = 0.02
project.simulation.max_step = 0.01
xml = build_reactflow_system_xml(project)
simulation_id = f"test-{uuid4().hex}"
task = _register_simulation_task(simulation_id)
events: list[dict[str, object]] = []
stop_requested = False
for line in simulation_event_stream(xml, task=task):
event = json.loads(line)
events.append(event)
if (
not stop_requested
and event["event"] == "progress"
and event["phase"] == "integrating"
and event["progress"] >= 20
):
response = cancel_system_xml_simulation(
simulation_id,
SimulationCancellationPayload(reason="user"),
)
self.assertTrue(response["accepted"])
stop_requested = True
result_event = next(event for event in events if event["event"] == "result")
result = result_event["result"]
snapshot = get_system_xml_simulation(simulation_id)
self.assertTrue(stop_requested)
self.assertEqual(result["status"], "stopped")
self.assertTrue(result["partial"])
self.assertGreaterEqual(result["diagnostics"]["sampleCount"], 2)
self.assertGreater(result["simulatedUntil"], 0.0)
self.assertLess(result["simulatedUntil"], 2.0)
self.assertEqual(snapshot["status"], "stopped")
self.assertEqual(snapshot["result"]["status"], "stopped")
def test_streaming_endpoint_keeps_quiet_solver_connection_alive(self) -> None:
def delayed_simulation(_xml_bytes, progress_callback, _cancel_check=None):
progress_callback(49, "integrating", "正在进行时间积分与压力流量求解")
time.sleep(0.03)
return {"success": True}
with (
patch("app.main.SIMULATION_STREAM_HEARTBEAT_SECONDS", 0.005),
patch(
"app.main.run_system_xml_simulation",
side_effect=delayed_simulation,
),
):
events = [
json.loads(line)
for line in simulation_event_stream(b"")
]
heartbeats = [event for event in events if event.get("heartbeat") is True]
self.assertGreaterEqual(len(heartbeats), 1)
self.assertTrue(all(event["progress"] == 49 for event in heartbeats))
self.assertTrue(all(event["phase"] == "integrating" for event in heartbeats))
self.assertEqual(events[-1]["event"], "result")
def test_streaming_endpoint_returns_structured_validation_error(self) -> None:
events = [
json.loads(line)
for line in simulation_event_stream(b"")
]
self.assertEqual(events[-1]["event"], "error")
self.assertEqual(events[-1]["status"], 422)
self.assertIn("issues", events[-1]["detail"])
def test_simulation_endpoint_returns_422_for_ideal_storage_coupling(self) -> None:
project = chain_project()
project.nodes = [project.nodes[0], project.nodes[-1]]
project.edges = [
ReactFlowEdgePayload(
**physical_edge(
"edge-1",
"cylinder_1",
"port_b",
"tank_1",
"port_a",
)
)
]
with self.assertRaises(HTTPException) as caught:
asyncio.run(
simulate_system_xml(
xml_request(build_reactflow_system_xml(project))
)
)
self.assertEqual(caught.exception.status_code, 422)
self.assertIn(
"IDEAL_STORAGE_COUPLING_UNSUPPORTED",
{
issue["code"]
for issue in caught.exception.detail["issues"]
},
)
if __name__ == "__main__":
unittest.main()