diff --git a/pslab/cli.py b/pslab/cli.py index 5ba833b..241b0a6 100644 --- a/pslab/cli.py +++ b/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): diff --git a/pslab/sciencelab.py b/pslab/sciencelab.py index 1cbc81e..cd52e51 100644 --- a/pslab/sciencelab.py +++ b/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. diff --git a/tests/test_flash.py b/tests/test_flash.py new file mode 100644 index 0000000..d9aebbb --- /dev/null +++ b/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"