fix: make pslab flash work again after the SerialHandler refactor

pslab/cli.py

@@ -237,7 +237,7 @@ def main(args: argparse.Namespace):
     handler = SerialHandler(port=args.port)
 
     if args.function == "flash":
-        flash(pslab.ScienceLab(args.port), args.hexfile)
+        flash(pslab.ScienceLab(handler), args.hexfile)
         return
 
     if args.function == "collect":
@@ -530,29 +530,31 @@ def flash(psl: pslab.ScienceLab, hexfile: str):
 
     PSLab must be in bootloader mode.
     """
-    if psl.interface.baudrate == 1000000:
-        psl.interface.timeout = 5
+    device = psl.device
+
+    if device.baudrate == 1000000:
+        device.timeout = 5
         psl.enter_bootloader()
 
     try:
-        bootattrs = mcbootflash.get_boot_attrs(psl)
+        bootattrs = mcbootflash.get_boot_attrs(device)
     except struct.error:
         print("Flashing failed: PSLab is not in bootloader mode.")
         return
 
-    mcbootflash.erase_flash(psl, bootattrs.memory_range, bootattrs.erase_size)
+    mcbootflash.erase_flash(device, bootattrs.memory_range, bootattrs.erase_size)
     total_bytes, chunks = mcbootflash.chunked(hexfile, bootattrs)
     written = 0
 
     for chunk in chunks:
-        mcbootflash.write_flash(psl, chunk)
-        mcbootflash.checksum(psl, chunk)
+        mcbootflash.write_flash(device, chunk)
+        mcbootflash.checksum(device, chunk)
         written += len(chunk.data)
         print(f"{written}/{total_bytes} bytes flashed.", end="\r")
 
     print("", end="\n")
-    mcbootflash.self_verify(psl)
-    mcbootflash.reset(psl)
+    mcbootflash.self_verify(device)
+    mcbootflash.reset(device)
 
 
 def add_flash_args(subparser: argparse._SubParsersAction):

pslab/sciencelab.py

@@ -123,7 +123,7 @@ def enter_bootloader(self):
             raise RuntimeError(msg)
 
         self.reset()
-        self.device.interface.baudrate = 460800
+        self.device.baudrate = 460800
         # The PSLab's RGB LED flashes some colors on boot.
         boot_lightshow_time = 0.6
         # Wait before sending magic number to make sure UART is initialized.

tests/test_flash.py

@@ -0,0 +1,68 @@
+"""Tests for flashing firmware from the CLI.
+
+These tests do not require a connected PSLab.
+"""
+
+import argparse
+from unittest.mock import MagicMock
+
+import pytest
+import serial
+
+import pslab
+from pslab import cli
+from pslab.connection import SerialHandler
+
+
+@pytest.fixture
+def handler() -> SerialHandler:
+    sh = SerialHandler("loop://")
+    sh._ser = serial.serial_for_url("loop://", baudrate=1000000, timeout=0.1)
+    yield sh
+    sh._ser.close()
+
+
+@pytest.fixture
+def psl(handler: SerialHandler) -> pslab.ScienceLab:
+    psl = pslab.ScienceLab.__new__(pslab.ScienceLab)
+    psl.device = handler
+    return psl
+
+
+def test_enter_bootloader_switches_baudrate(psl, monkeypatch):
+    monkeypatch.setattr("time.sleep", lambda _: None)
+
+    psl.enter_bootloader()
+
+    assert psl.device.baudrate == 460800
+
+
+def test_flash_uses_the_serial_connection(psl, monkeypatch):
+    bootflash = MagicMock()
+    bootflash.chunked.return_value = (0, [])
+    monkeypatch.setattr(cli, "mcbootflash", bootflash)
+    monkeypatch.setattr(psl, "enter_bootloader", MagicMock())
+
+    cli.flash(psl, "firmware.hex")
+
+    psl.enter_bootloader.assert_called_once()
+    assert psl.device.timeout == 5
+    bootflash.get_boot_attrs.assert_called_once_with(psl.device)
+    bootflash.reset.assert_called_once_with(psl.device)
+
+
+def test_main_flashes_through_the_opened_handler(handler, monkeypatch):
+    monkeypatch.setattr(cli, "SerialHandler", lambda port: handler)
+    monkeypatch.setattr(
+        pslab.ScienceLab,
+        "__init__",
+        lambda self, device: setattr(self, "device", device),
+    )
+    flash = MagicMock()
+    monkeypatch.setattr(cli, "flash", flash)
+
+    cli.main(argparse.Namespace(function="flash", port=None, hexfile="firmware.hex"))
+
+    psl, hexfile = flash.call_args.args
+    assert psl.device is handler
+    assert hexfile == "firmware.hex"