feat: add Modbus drivers for Jig, TPU, and RPU

drivers/jig.py

@@ -0,0 +1,40 @@
+import time
+from dataclasses import dataclass
+
+from drivers.modbus_device import ModbusDevice, ModbusDeviceError
+
+
+@dataclass
+class Position:
+    x: int
+    y: int
+    z: int
+
+    def __str__(self):
+        return f"(x={self.x}, y={self.y}, z={self.z})mm"
+
+
+class JigMovementError(ModbusDeviceError):
+    pass
+
+
+class Jig(ModbusDevice):
+    def __init__(self, port, registers, slave_id=0x10):
+        super().__init__(port, "Jig", registers, slave_id=slave_id)
+
+    def move_to_position(self, position, wait_time=5):
+        try:
+            self.write_holding_register(self._addr("X POSITION"), position.x)
+            self.write_holding_register(self._addr("Y POSITION"), position.y)
+            self.write_holding_register(self._addr("Z POSITION"), position.z)
+            self.write_coil(self._addr("MOVE TO POSITION"), True, verify=False)
+            if wait_time > 0:
+                time.sleep(wait_time)
+        except ModbusDeviceError as exc:
+            raise JigMovementError(f"failed to move to {position}") from exc
+
+    def get_current_position(self):
+        x = self.read_holding_register(self._addr("X POSITION"))
+        y = self.read_holding_register(self._addr("Y POSITION"))
+        z = self.read_holding_register(self._addr("Z POSITION"))
+        return Position(x, y, z)

drivers/modbus_device.py

@@ -0,0 +1,186 @@
+import struct
+import time
+
+from pymodbus.client import ModbusSerialClient
+from pymodbus.exceptions import ModbusException
+
+
+class ModbusDeviceError(Exception):
+    pass
+
+
+class ModbusConnectionError(ModbusDeviceError):
+    pass
+
+
+class ModbusReadError(ModbusDeviceError):
+    pass
+
+
+class ModbusWriteError(ModbusDeviceError):
+    pass
+
+
+class ModbusDevice:
+    FLOAT_VERIFY_TOLERANCE = 0.01
+
+    def __init__(self, port, device_name, registers=None, slave_id=0x10, baudrate=9600, timeout=3,
+                 retries=3, retry_delay_s=0.2):
+        self.port = port
+        self.device_name = device_name
+        self.registers = registers or {}
+        self.slave_id = slave_id
+        self.retries = retries
+        self.retry_delay_s = retry_delay_s
+        try:
+            self.client = ModbusSerialClient(port=port, baudrate=baudrate, parity="N", stopbits=1, bytesize=8, timeout=timeout)
+            if not self.client.connect():
+                raise ModbusConnectionError(f"failed to connect to {device_name} on {port}")
+        except ModbusException as exc:
+            raise ModbusConnectionError(f"error connecting to {device_name} on {port}") from exc
+
+    def _addr(self, label):
+        try:
+            return self.registers[label].address
+        except KeyError:
+            raise ModbusDeviceError(f"unknown register label for {self.device_name}: {label}")
+
+    def _retry(self, action):
+        last_exc = None
+        for attempt in range(self.retries):
+            try:
+                return action()
+            except (ModbusDeviceError, ModbusException) as exc:
+                last_exc = exc
+                if attempt < self.retries - 1:
+                    time.sleep(self.retry_delay_s)
+        raise last_exc
+
+    # ---- raw reads (retried on failure) ----
+
+    def read_coil(self, address):
+        def _do():
+            res = self.client.read_coils(address, count=1, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusReadError(f"{self.device_name}: failed to read coil {address}")
+            return res.bits[0]
+        return self._retry(_do)
+
+    def read_holding_register(self, address):
+        def _do():
+            res = self.client.read_holding_registers(address, count=1, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusReadError(f"{self.device_name}: failed to read holding {address}")
+            return res.registers[0]
+        return self._retry(_do)
+
+    def read_holding_register_float(self, address):
+        def _do():
+            res = self.client.read_holding_registers(address, count=2, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusReadError(f"{self.device_name}: failed to read holding float {address}")
+            return struct.unpack(">f", struct.pack(">HH", res.registers[0], res.registers[1]))[0]
+        return self._retry(_do)
+
+    def read_input_register(self, address):
+        def _do():
+            res = self.client.read_input_registers(address, count=1, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusReadError(f"{self.device_name}: failed to read input {address}")
+            return res.registers[0]
+        return self._retry(_do)
+
+    def read_signed_input_register(self, address):
+        val = self.read_input_register(address)
+        return val - 65536 if val > 32767 else val
+
+    def read_input_register_float(self, address):
+        def _do():
+            res = self.client.read_input_registers(address, count=2, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusReadError(f"{self.device_name}: failed to read input float {address}")
+            return struct.unpack(">f", struct.pack(">HH", res.registers[0], res.registers[1]))[0]
+        return self._retry(_do)
+
+    # ---- writes, verified by reading the value back (whole write+verify retried on mismatch) ----
+
+    def write_coil(self, address, value, verify=True):
+        value = bool(value)
+
+        def _do():
+            res = self.client.write_coil(address, value, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusWriteError(f"{self.device_name}: failed to write coil {address}")
+            if verify and self.read_coil(address) != value:
+                raise ModbusWriteError(f"{self.device_name}: write coil {address}={value} not verified")
+        self._retry(_do)
+
+    def write_holding_register(self, address, value):
+        value = int(value) & 0xFFFF  # two's-complement wrap so negative (signed) values fit the uint16 wire format
+
+        def _do():
+            res = self.client.write_register(address, value, device_id=self.slave_id)
+            if res.isError():
+                raise ModbusWriteError(f"{self.device_name}: failed to write holding {address}")
+            if self.read_holding_register(address) != value:
+                raise ModbusWriteError(f"{self.device_name}: write holding {address}={value} not verified")
+        self._retry(_do)
+
+    def read_signed_holding_register(self, address):
+        val = self.read_holding_register(address)
+        return val - 65536 if val > 32767 else val
+
+    def write_holding_register_float(self, address, value):
+        def _do():
+            reg1, reg2 = struct.unpack(">HH", struct.pack(">f", value))
+            res = self.client.write_registers(address, [reg1, reg2], device_id=self.slave_id)
+            if res.isError():
+                raise ModbusWriteError(f"{self.device_name}: failed to write holding float {address}")
+            readback = self.read_holding_register_float(address)
+            if abs(readback - value) > self.FLOAT_VERIFY_TOLERANCE:
+                raise ModbusWriteError(
+                    f"{self.device_name}: write holding float {address}={value} not verified (read back {readback})"
+                )
+        self._retry(_do)
+
+    # ---- generic label-based access (covers every register in self.registers) ----
+
+    def _register(self, label):
+        try:
+            return self.registers[label]
+        except KeyError:
+            raise ModbusDeviceError(f"unknown register label for {self.device_name}: {label}")
+
+    def read(self, label):
+        reg = self._register(label)
+        if reg.object_type == "coil":
+            return self.read_coil(reg.address)
+        if reg.object_type == "input":
+            if reg.data_type == "float":
+                return self.read_input_register_float(reg.address)
+            if reg.data_type == "int":
+                return self.read_signed_input_register(reg.address)
+            return self.read_input_register(reg.address)
+        if reg.object_type == "holding":
+            if reg.data_type == "float":
+                return self.read_holding_register_float(reg.address)
+            if reg.data_type == "int":
+                return self.read_signed_holding_register(reg.address)
+            return self.read_holding_register(reg.address)
+        raise ModbusDeviceError(f"{self.device_name}: unsupported object_type {reg.object_type!r} for {label!r}")
+
+    def write(self, label, value):
+        reg = self._register(label)
+        if reg.object_type == "coil":
+            self.write_coil(reg.address, value)
+            return
+        if reg.object_type == "holding":
+            if reg.data_type == "float":
+                self.write_holding_register_float(reg.address, value)
+            else:
+                self.write_holding_register(reg.address, value)
+            return
+        raise ModbusDeviceError(f"{self.device_name}: cannot write to object_type {reg.object_type!r} for {label!r}")
+
+    def close(self):
+        self.client.close()

drivers/rpu.py

@@ -0,0 +1,43 @@
+import time
+
+from common.config_loader import load_garuda_registers
+from drivers.modbus_device import ModbusDevice
+
+DEFAULT_DATASETS = ["Sec_normal", "sec_flash"]
+
+
+class Rpu(ModbusDevice):
+    """Covers every register in Sec_normal + sec_flash by label via read()/write()
+    (inherited from ModbusDevice). Extra datasets can be added per-test via
+    extra_datasets. Only registers that need real logic beyond a 1:1 read/write
+    get a dedicated method below.
+    """
+
+    Q128_DIVISOR = 128.0
+    REWRITE_PASSWORD_LO = 0xFEED
+    REWRITE_PASSWORD_HI = 0xCEE5
+    REWRITE_WAIT_S = 5
+
+    def __init__(self, port, extra_datasets=None, slave_id=0x10):
+        registers = load_garuda_registers(DEFAULT_DATASETS + list(extra_datasets or []))
+        super().__init__(port, "RPU", registers, slave_id=slave_id)
+
+    def set_charging(self, enable):
+        self.write_coil(self._addr("Enable_charger"), enable)
+
+    def set_runtime_voltage(self, volts):
+        self.write_holding_register(self._addr("Runtime_voltage_setting_Q128"), round(volts * self.Q128_DIVISOR))
+
+    def set_runtime_current(self, amps):
+        self.write_holding_register(self._addr("Runtime_current_setting_Q128"), round(amps * self.Q128_DIVISOR))
+
+    def read_error_code(self):
+        return (self.read("Error_code_hi") << 16) | self.read("Error_code_lo")
+
+    def read_shadow_error_code(self):
+        return (self.read("Shadow_error_code_hi") << 16) | self.read("Shadow_error_code_lo")
+
+    def rewrite_config(self):
+        self.write_holding_register(self._addr("PASSWORD_LO"), self.REWRITE_PASSWORD_LO)
+        self.write_holding_register(self._addr("PASSWORD_HI"), self.REWRITE_PASSWORD_HI)
+        time.sleep(self.REWRITE_WAIT_S)

drivers/tpu.py

@@ -0,0 +1,44 @@
+import time
+
+from common.config_loader import load_fenghuang_registers
+from drivers.modbus_device import ModbusDevice, ModbusDeviceError
+
+DEFAULT_DATASETS = ["Pri_normal", "Debug", "event_log"]
+
+# TPU_STATE values that allow modbus_to_config()/write_all_flash() to run
+# (fenghuang-dsp/app/app_modbus/app_modbus.c update_modbus_input()).
+# TPU_RUN_INV (charging) is NOT in this set - config writes are silently
+# ignored while charging is active.
+CONFIG_WRITABLE_STATES = {2, 10, 15, 16, 238}  # READY, INITIALIZATION, DEBUG_CONFIG, TEST_MODE, ERROR
+
+
+class Tpu(ModbusDevice):
+    """Covers every register in Pri_normal + Debug + event_log by label via
+    read()/write() (inherited from ModbusDevice). More datasets can be added
+    per-test via extra_datasets. Only registers that need real logic beyond a
+    1:1 read/write get a dedicated method below.
+    """
+
+    REWRITE_WAIT_S = 5
+
+    def __init__(self, port, extra_datasets=None, slave_id=0x10):
+        registers = load_fenghuang_registers(DEFAULT_DATASETS + list(extra_datasets or []))
+        super().__init__(port, "TPU", registers, slave_id=slave_id)
+
+    def wait_until_config_writable(self, timeout_s=10, poll_interval_s=0.5):
+        deadline = time.time() + timeout_s
+        while time.time() < deadline:
+            if self.read("Main_state") in CONFIG_WRITABLE_STATES:
+                return
+            time.sleep(poll_interval_s)
+        raise ModbusDeviceError(f"TPU did not reach a config-writable state within {timeout_s}s")
+
+    def read_error_code(self):
+        return (self.read("Error_code_HI") << 16) | self.read("Error_code_LO")
+
+    def read_shadow_error_code(self):
+        return (self.read("Shadow_Error_code_HI") << 16) | self.read("Shadow_Error_code_LO")
+
+    def rewrite_config(self):
+        self.write_coil(self._addr("rewrite_cfg_flash"), True, verify=False)
+        time.sleep(self.REWRITE_WAIT_S)