Add temperature tracking (MOS and rotor) to motor state

src/openarm_driver/driver.py

@@ -132,14 +132,17 @@ def set_latest_state(self, timeout_us=300):
         """Update the state."""
         self.openarm.recv_all(timeout_us)
         motor_values = (
-            (m.get_position(), m.get_velocity(), m.get_torque())
+            (m.get_position(), m.get_velocity(), m.get_torque(),
+             m.get_state_tmos(), m.get_state_trotor())
             for m in self._iter_motors()
         )
-        qpos, qvel, qtau = zip(*motor_values)
+        qpos, qvel, qtau, tmos, trotor = zip(*motor_values)
         self.latest_state = {
             "qpos": np.array(qpos, dtype=float) - self.joint_offsets,
             "qvel": np.array(qvel, dtype=float),
             "qtorque": np.array(qtau, dtype=float),
+            "tmos": np.array(tmos, dtype=int),
+            "trotor": np.array(trotor, dtype=int),
         }
 
     def fetch_state(self, refresh=True) -> dict[str, np.ndarray]:
@@ -162,6 +165,14 @@ def fetch_torque(self, refresh=True) -> np.ndarray:
         """Fetch the torque."""
         return self.fetch_state(refresh=refresh)["qtorque"]
 
+    def fetch_mos_temperature(self, refresh=True) -> np.ndarray:
+        """Fetch the MOS temperature for each motor."""
+        return self.fetch_state(refresh=refresh)["tmos"]
+
+    def fetch_rotor_temperature(self, refresh=True) -> np.ndarray:
+        """Fetch the rotor temperature for each motor."""
+        return self.fetch_state(refresh=refresh)["trotor"]
+
     def send_position(self, position: ArrayLike):
         """Move the arm by sending the position."""
         checked_result = self.safety_checker.check(position, driver=self)

tests/test_driver.py

@@ -22,6 +22,8 @@ def __init__(self):
         self.position = 0.5
         self.velocity = 0.0
         self.torque = 0.0
+        self.tmos = 25
+        self.trotor = 30
 
     def get_position(self):
         return self.position
@@ -32,6 +34,12 @@ def get_velocity(self):
     def get_torque(self):
         return self.torque
 
+    def get_state_tmos(self):
+        return self.tmos
+
+    def get_state_trotor(self):
+        return self.trotor
+
 
 class CanMock:
     def __init__(self, *args, **kwargs):
@@ -95,6 +103,20 @@ def test_fetch_torque(can_mock):
     driver.fetch_torque(refresh=False)
 
 
+def test_fetch_mos_temperature(can_mock):
+    driver = SingleArmDriver("right_arm")
+    temps = driver.fetch_mos_temperature(refresh=True)
+    assert temps.tolist() == [25] * 8
+    driver.fetch_mos_temperature(refresh=False)
+
+
+def test_fetch_rotor_temperature(can_mock):
+    driver = SingleArmDriver("right_arm")
+    temps = driver.fetch_rotor_temperature(refresh=True)
+    assert temps.tolist() == [30] * 8
+    driver.fetch_rotor_temperature(refresh=False)
+
+
 def test_fetch_state(can_mock):
     driver = SingleArmDriver("right_arm")
     driver.fetch_state(refresh=True)