from __future__ import annotations import asyncio import unittest from fastapi import HTTPException, Request from app.main import ( ReactFlowEdgePayload, ReactFlowProjectPayload, build_reactflow_system_xml, compile_reactflow_network, 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()) result = GenericFluidSystem(network).simulate( SolveIVPConfig( t_start=0.0, t_stop=0.01, method="BDF", max_step=0.001, ), sample_step=0.005, ) self.assertTrue(result.success) 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_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) 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()