From d4edf1ccb33f162be792402bb761ba85c9fbbf1f Mon Sep 17 00:00:00 2001
From: huojiarong <huojiarong@sf.banytech.com>
Date: Mon, 20 Jul 2026 09:26:01 +0000
Subject: [PATCH] =?UTF-8?q?=E5=AF=BC=E5=87=BAtest=5Fmql=E6=B4=BB=E5=A1=9E?=
 =?UTF-8?q?=E8=BF=90=E5=8A=A8=E5=AD=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 PythonModels/systems/test_mql_mechanical.py | 45 ++++++++++++++++++++-
 tests/test_test_mql_mechanical.py           | 24 +++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/PythonModels/systems/test_mql_mechanical.py b/PythonModels/systems/test_mql_mechanical.py
index 4591b3c..1b8481f 100644
--- a/PythonModels/systems/test_mql_mechanical.py
+++ b/PythonModels/systems/test_mql_mechanical.py
@@ -210,10 +210,20 @@ class TestMqlMechanicalNodeKinematics:
     displacements_m: dict[int, float]
 
 
+@dataclass(frozen=True)
+class TestMqlPistonKinematics:
+    alias: str
+    port_2_velocity_m_s: float
+    port_2_displacement_m: float
+    port_3_velocity_m_s: float
+    port_3_displacement_m: float
+
+
 @dataclass(frozen=True)
 class TestMqlMechanicalMassSnapshot:
     states: tuple[TestMqlMechanicalMassState, ...]
     node_kinematics_by_alias: dict[str, TestMqlMechanicalNodeKinematics]
+    piston_kinematics_by_alias: dict[str, TestMqlPistonKinematics]
 
     @property
     def state_count(self) -> int:
@@ -244,9 +254,14 @@ class TestMqlMechanicalMassClosure:
             )
             for index, alias in enumerate(self.mass_aliases)
         )
+        node_kinematics = self._node_kinematics_by_alias(states)
         return TestMqlMechanicalMassSnapshot(
             states=states,
-            node_kinematics_by_alias=self._node_kinematics_by_alias(states),
+            node_kinematics_by_alias=node_kinematics,
+            piston_kinematics_by_alias=self._piston_kinematics_by_alias(
+                states,
+                node_kinematics,
+            ),
         )
 
     def _node_kinematics_by_alias(
@@ -269,6 +284,34 @@ class TestMqlMechanicalMassClosure:
             ),
         }
 
+    def _piston_kinematics_by_alias(
+        self,
+        states: tuple[TestMqlMechanicalMassState, ...],
+        node_kinematics_by_alias: dict[str, TestMqlMechanicalNodeKinematics],
+    ) -> dict[str, TestMqlPistonKinematics]:
+        state_by_alias = {state.alias: state for state in states}
+        rear_node = node_kinematics_by_alias["dynamic_mechanical_node_alternative_3"]
+        piston_bindings = (
+            ("pn_brp2_8", "mass_friction_endstops_10", 8),
+            ("pn_brp2_9", "mass_friction_endstops_11", 7),
+            ("pn_brp2_10", "mass_friction_endstops_12", 6),
+            ("pn_brp2_11", "mass_friction_endstops_13", 5),
+            ("pn_brp2_12", "mass_friction_endstops_14", 4),
+            ("pn_brp2_13", "mass_friction_endstops_15", 3),
+            ("pn_brp2_14", "mass_friction_endstops_16", 2),
+            ("pn_brp2_15", "mass_friction_endstops_17", 1),
+        )
+        return {
+            piston_alias: TestMqlPistonKinematics(
+                alias=piston_alias,
+                port_2_velocity_m_s=state_by_alias[mass_alias].velocity_m_s,
+                port_2_displacement_m=state_by_alias[mass_alias].displacement_m,
+                port_3_velocity_m_s=rear_node.velocities_m_s[rear_node_port],
+                port_3_displacement_m=rear_node.displacements_m[rear_node_port],
+            )
+            for piston_alias, mass_alias, rear_node_port in piston_bindings
+        }
+
     def rhs(
         self,
         state_vector: list[float],
diff --git a/tests/test_test_mql_mechanical.py b/tests/test_test_mql_mechanical.py
index b6c4681..3d74536 100644
--- a/tests/test_test_mql_mechanical.py
+++ b/tests/test_test_mql_mechanical.py
@@ -120,6 +120,30 @@ class TestMqlMechanicalAssemblyTests(unittest.TestCase):
         self.assertAlmostEqual(rear.velocities_m_s[1], -0.3)
         self.assertAlmostEqual(rear.displacements_m[8], -0.72)
 
+    def test_mechanical_mass_snapshot_exposes_piston_kinematics(self) -> None:
+        closure = TestMqlMechanicalMassClosure(self.assembly)
+        state = closure.initial_state_vector()
+        state[0] = 0.4
+        state[1] = 0.37
+        state[18] = -0.3
+        state[19] = -0.72
+
+        snapshot = closure.snapshot(state)
+        piston = snapshot.piston_kinematics_by_alias["pn_brp2_8"]
+
+        self.assertEqual(set(snapshot.piston_kinematics_by_alias), set(self.assembly.pistons))
+        self.assertAlmostEqual(piston.port_2_velocity_m_s, 0.4)
+        self.assertAlmostEqual(piston.port_2_displacement_m, 0.37)
+        self.assertAlmostEqual(piston.port_3_velocity_m_s, -0.3)
+        self.assertAlmostEqual(piston.port_3_displacement_m, -0.72)
+        self.assertAlmostEqual(
+            self.assembly.pistons["pn_brp2_8"].geometry().chamber_length_mm(
+                piston.port_3_displacement_m,
+                piston.port_2_displacement_m,
+            ),
+            1090.0,
+        )
+
     def test_system_simulates_mechanical_mass_state_closure(self) -> None:
         system = TestMqlSystem()