diff --git a/src/openarm_driver/driver.py b/src/openarm_driver/driver.py index 023136a..6af12b2 100644 --- a/src/openarm_driver/driver.py +++ b/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) diff --git a/tests/test_driver.py b/tests/test_driver.py index b5095fd..24d2ef7 100644 --- a/tests/test_driver.py +++ b/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)