feat(heater-shaker): add firmware upload capability (#10307)

* added upload_via_dfu method and updated find_bootloader_port for stm32 modules

* formatted

* fix for testing

* fix for testing

* Made PR requested changes and added test

* formatted

* fixes from linting

* updated tests

* added test

* updated tests

* updated the find-and-enter-bootloader flow for h/s

* don't use returncode for long processes

* linted, updated tests, and made some suggested changes

* halp

* formatted

* updated locking mechanism

* removed comments

Co-authored-by: Sanniti <[email protected]>

api/src/opentrons/drivers/asyncio/communication/serial_connection.py

@@ -119,6 +119,19 @@ async def send_command(
  data=command.build(), retries=retries, timeout=timeout
  )
 
+ async def send_dfu_command(self, command: CommandBuilder) -> None:
+ """
+ Send a dfu command to enter device bootloader.
+
+ This method doesn't wait for a response after sending the command since the
+ module port gets disconnected once it enters its bootloader.
+ """
+ encoded_command = command.build().encode()
+
+ async with self._send_data_lock:
+ log.debug(f"{self.name}: Write -> {encoded_command!r}")
+ await self._serial.write(data=encoded_command)
+
  async def send_data(
  self, data: str, retries: int = 0, timeout: Optional[float] = None
  ) -> str:

api/src/opentrons/drivers/heater_shaker/driver.py

@@ -176,7 +176,7 @@ async def enter_programming_mode(self) -> None:
  c = CommandBuilder(terminator=HS_COMMAND_TERMINATOR).add_gcode(
  gcode=GCODE.ENTER_BOOTLOADER
  )
- await self._connection.send_command(command=c, retries=DEFAULT_COMMAND_RETRIES)
+ await self._connection.send_dfu_command(command=c)
  await self._connection.close()
 
  async def deactivate_heater(self) -> None:

api/src/opentrons/hardware_control/modules/heater_shaker.py

@@ -29,6 +29,8 @@
 # module simulation in PAPIv2 needs to be seriously rethought
 SIMULATING_POLL_PERIOD = POLL_PERIOD / 20.0
 
+DFU_PID = "df11"
+
 
 class HeaterShakerError(RuntimeError):
  """An error propagated from the heater-shaker module."""
@@ -439,7 +441,11 @@ async def close_labware_latch(self) -> None:
  await self._wait_for_labware_latch(HeaterShakerLabwareLatchStatus.IDLE_CLOSED)
 
  async def prep_for_update(self) -> str:
- return "no"
+ await self._poller.stop_and_wait()
+ async with update.protect_dfu_transition():
+ await self._driver.enter_programming_mode()
+ dfu_info = await update.find_dfu_device(pid=DFU_PID)
+ return dfu_info
 
 
 @dataclass

api/src/opentrons/hardware_control/modules/mod_abc.py

@@ -76,6 +76,7 @@ def get_bundled_fw(self) -> Optional[BundledFirmware]:
  name_to_fw_file_prefix = {
  "tempdeck": "temperature-module",
  "magdeck": "magnetic-module",
+ "heatershaker": "heater-shaker",
  }
  name = self.name()
  file_prefix = name_to_fw_file_prefix.get(name, name)

api/src/opentrons/hardware_control/modules/update.py

@@ -3,12 +3,22 @@
 import os
 from pathlib import Path
 from glob import glob
-from typing import Any, Dict, Tuple, Optional, Union
+from typing import Any, AsyncGenerator, Dict, Tuple, Optional, Union
 from .types import UpdateError
 from .mod_abc import AbstractModule
+from opentrons.hardware_control.threaded_async_lock import ThreadedAsyncLock
+from contextlib import asynccontextmanager
 
 log = logging.getLogger(__name__)
 
+_dfu_transition_lock = ThreadedAsyncLock()
+
+
+@asynccontextmanager
+async def protect_dfu_transition() -> AsyncGenerator[None, None]:
+ async with _dfu_transition_lock.lock():
+ yield
+
 
 async def update_firmware(
  module: AbstractModule,
@@ -19,13 +29,15 @@ async def update_firmware(
 
  raises an UpdateError with the reason for the failure.
  """
- flash_port = await module.prep_for_update()
+ flash_port_or_dfu_serial = await module.prep_for_update()
  kwargs: Dict[str, Any] = {
  "stdout": asyncio.subprocess.PIPE,
  "stderr": asyncio.subprocess.PIPE,
  "loop": loop,
  }
- successful, res = await module.bootloader()(flash_port, str(firmware_file), kwargs)
+ successful, res = await module.bootloader()(
+ flash_port_or_dfu_serial, str(firmware_file), kwargs
+ )
  if not successful:
  log.info(f"Bootloader reponse: {res}")
  raise UpdateError(res)
@@ -50,6 +62,48 @@ async def find_bootloader_port() -> str:
  raise Exception("No ot_module bootloaders found in /dev. Try again")
 
 
+async def find_dfu_device(pid: str) -> str:
+ """Find the dfu device and return its serial number (separate from module serial)"""
+ retries = 5
+ log.info(f"Searching for a dfu device with PID {pid}")
+ while retries != 0:
+ retries -= 1
+ await asyncio.sleep(1)
+ proc = await asyncio.create_subprocess_exec(
+ "dfu-util",
+ "-l",
+ stdout=asyncio.subprocess.PIPE,
+ stderr=asyncio.subprocess.PIPE,
+ )
+ await proc.wait()
+ stdout, stderr = await proc.communicate()
+
+ if stdout is None and stderr is None:
+ continue
+ if stderr:
+ raise RuntimeError(f"Error finding dfu device: {stderr.decode()}")
+
+ result = stdout.decode()
+ if pid not in result:
+ # It could take a few seconds for the device to show up
+ continue
+ devices_found = 0
+ for line in result.splitlines():
+ if pid in line:
+ log.info(f"Found device with PID {pid}")
+ devices_found += 1
+ serial = line[(line.find("serial=") + 7) :]
+ if devices_found == 2:
+ return serial
+ elif devices_found > 2:
+ raise OSError("Multiple new bootloader devices" "found on mode switch")
+
+ raise RuntimeError(
+ "Could not update firmware via dfu. Possible issues- dfu-util"
+ " not working or specified dfu device not found"
+ )
+
+
 async def upload_via_avrdude(
  port: str, firmware_file_path: str, kwargs: Dict[str, Any]
 ) -> Tuple[bool, str]:
@@ -141,10 +195,37 @@ async def upload_via_bossa(
 
 
 async def upload_via_dfu(
- port: str, firmware_file: str, kwargs: Dict[str, Any]
+ dfu_serial: str, firmware_file_path: str, kwargs: Dict[str, Any]
 ) -> Tuple[bool, str]:
- """Run a TBD firmware upload command for heater/shaker.
+ """Run firmware upload command for DFU.
 
  Returns tuple of success boolean and message from bootloader
  """
- return False, ""
+ log.info("Starting firmware upload via dfu util")
+
+ # We don't specify a port since the dfu device doesn't get one &
+ # dfu-util doesn't need it either so I've replaced the 'port' arg with an arg to
+ # fetch dfu device serial. I haven't used it in the below command though.
+ dfu_args = [
+ "dfu-util",
+ "-a 0",
+ "-s 0x08000000:leave",
+ f"-D{firmware_file_path}",
+ "-R",
+ ]
+ proc = await asyncio.create_subprocess_exec(*dfu_args, **kwargs)
+ stdout, stderr = await proc.communicate()
+ res = stdout.decode()
+
+ if "File downloaded successfully" in res:
+ log.debug(res)
+ log.info("Firmware upload successful")
+ return True, res
+ else:
+ log.error(
+ f"Failed to update module firmware for {dfu_serial}. "
+ # It isn't easy to decipher the issue from stderror alone
+ f"stdout: {res} \n"
+ f"stderr: {stderr.decode()}"
+ )
+ return False, res

api/tests/opentrons/drivers/heater_shaker/test_driver.py

@@ -145,10 +145,10 @@ async def test_enter_bootloader(
  subject: driver.HeaterShakerDriver, connection: AsyncMock
 ) -> None:
  """It should enter the bootloader"""
- connection.send_command.return_value = None
+ connection.send_dfu_command.return_value = None
  await subject.enter_programming_mode()
  expected = CommandBuilder(terminator=driver.HS_COMMAND_TERMINATOR).add_gcode(
  gcode="dfu"
  )
- connection.send_command.assert_called_once_with(command=expected, retries=0)
+ connection.send_dfu_command.assert_called_once_with(command=expected)
  connection.close.assert_called_once()

api/tests/opentrons/hardware_control/test_modules.py

@@ -243,8 +243,33 @@ async def mod_thermocycler():
  await thermocycler.cleanup()
 
 
+@pytest.fixture
+async def mod_heatershaker():
+ from opentrons.hardware_control import modules
+
+ loop = asyncio.get_running_loop()
+
+ usb_port = USBPort(
+ name="",
+ hub=None,
+ port_number=0,
+ device_path="/dev/ot_module_sim_heatershaker0",
+ )
+
+ heatershaker = await modules.build(
+ port="/dev/ot_module_sim_heatershaker0",
+ usb_port=usb_port,
+ which="heatershaker",
+ simulating=True,
+ loop=loop,
+ execution_manager=ExecutionManager(),
+ )
+ yield heatershaker
+ await heatershaker.cleanup()
+
+
 async def test_module_update_integration(
- monkeypatch, mod_tempdeck, mod_magdeck, mod_thermocycler
+ monkeypatch, mod_tempdeck, mod_magdeck, mod_thermocycler, mod_heatershaker
 ):
  from opentrons.hardware_control import modules
 
@@ -304,6 +329,22 @@ async def mock_find_bossa_bootloader_port():
  "ot_module_bossa_bootloader1", "fake_fw_file_path", bootloader_kwargs
  )
 
+ # test heater-shaker module update with dfu bootloader
+ upload_via_dfu_mock = mock.Mock(
+ return_value=(async_return((True, "dfu bootloader worked")))
+ )
+ monkeypatch.setattr(modules.update, "upload_via_dfu", upload_via_dfu_mock)
+
+ async def mock_find_dfu_device(pid: str):
+ return "df11"
+
+ monkeypatch.setattr(modules.update, "find_dfu_device", mock_find_dfu_device)
+
+ await modules.update_firmware(mod_heatershaker, "fake_fw_file_path", loop)
+ upload_via_dfu_mock.assert_called_once_with(
+ "df11", "fake_fw_file_path", bootloader_kwargs
+ )
+
 
 async def test_get_bundled_fw(monkeypatch, tmpdir):
  from opentrons.hardware_control import modules