from __future__ import annotations import asyncio import json import os 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 app.simulation.components.experimental.flow.resistive_pipe import ResistivePipe from app.simulation.registry import get_component_model_spec from app.simulation.solvers.solver import SolveIVPConfig from app.simulation.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, "modelVersion": get_component_model_spec(model_type).model_version, "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( projectSchemaVersion=1, 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( projectSchemaVersion=1, 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) pressure_flow = result.diagnostics["pressureFlow"] self.assertLess(pressure_flow["maxScaledResidual"], 1e-7) self.assertEqual( pressure_flow["solveCount"], pressure_flow["seededSolveCount"] + pressure_flow["nonlinearSolveCount"], ) self.assertAlmostEqual( pressure_flow["fastPathHitRate"], pressure_flow["seededSolveCount"] / pressure_flow["solveCount"], ) self.assertGreaterEqual( pressure_flow["residualEvaluationCount"], pressure_flow["optimizerEvaluationCount"], ) 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() with patch.dict( os.environ, {"SIMULATION_ODE_JACOBIAN_MODE": "optimized"}, ): 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) integration = result.diagnostics["integration"] self.assertEqual(integration["method"], "BDF") self.assertEqual(integration["segmentCount"], len(integration["segments"])) self.assertEqual( integration["totals"]["nfev"], sum(segment["nfev"] for segment in integration["segments"]), ) sparsity = integration["jacobianSparsity"] self.assertGreater(sparsity["nonzeroCount"], 0) self.assertGreater(sparsity["colorGroupCount"], 0) self.assertEqual( integration["jacobian"]["mode"], "scipySparseFiniteDifference", ) self.assertEqual( integration["jacobian"]["fallbackReason"], "denseStateDependencyPattern", ) self.assertEqual( integration["totals"]["finiteDifferenceRhsEstimate"], sum( segment["finiteDifferenceRhsEstimate"] for segment in integration["segments"] ), ) 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_raw_system_xml_serializes_more_than_legacy_sample_limit(self) -> None: project = chain_project() project.simulation.step = 0.0000005 xml = build_reactflow_system_xml(project) response = asyncio.run(simulate_system_xml(xml_request(xml))) round_tripped = json.loads(json.dumps(response)) self.assertTrue(round_tripped["success"]) self.assertEqual(round_tripped["diagnostics"]["sampleCount"], 20001) self.assertEqual(len(round_tripped["series"]["time"]), 20001) self.assertEqual(round_tripped["series"]["time"][0], 0.0) self.assertEqual(round_tripped["series"]["time"][-1], 0.01) 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()