feat(arduino): Add support for fast flashing with esptool

Author: lucasssvaz

pytest-embedded-arduino/pytest_embedded_arduino/app.py

@@ -1,9 +1,12 @@
 import json
 import logging
 import os
+import re
 
 from pytest_embedded.app import App
 
+_HEX_ADDR_RE = re.compile(r'^0x[0-9a-fA-F]+$')
+
 
 class ArduinoApp(App):
     """
@@ -14,6 +17,8 @@ class ArduinoApp(App):
         fqbn (str): Fully Qualified Board Name.
         target (str) : ESPxx chip.
         flash_settings (dict[str, str]): Flash settings for the target.
+        flash_files (list[tuple[str, str]]): ``(address, filepath)`` pairs parsed
+            from ``flash_args``.  Each filepath is absolute.
         binary_file (str): Merged binary file path.
         elf_file (str): ELF file path.
     """
@@ -29,7 +34,7 @@ def __init__(
         self.sketch = self._get_sketch_name(self.binary_path)
         self.fqbn = self._get_fqbn(self.binary_path)
         self.target = self.fqbn.split(':')[2]
-        self.flash_settings = self._get_flash_settings()
+        self.flash_settings, self.flash_files = self._parse_flash_args()
         self.binary_file = os.path.realpath(os.path.join(self.binary_path, self.sketch + '.ino.merged.bin'))
         self.elf_file = os.path.realpath(os.path.join(self.binary_path, self.sketch + '.ino.elf'))
 
@@ -38,6 +43,7 @@ def __init__(
         logging.debug(f'FQBN: {self.fqbn}')
         logging.debug(f'Target: {self.target}')
         logging.debug(f'Flash settings: {self.flash_settings}')
+        logging.debug(f'Flash files: {self.flash_files}')
         logging.debug(f'Binary file: {self.binary_file}')
         logging.debug(f'ELF file: {self.elf_file}')
 
@@ -90,18 +96,45 @@ def _get_fqbn(self, build_path: str) -> str:
         fqbn = options['fqbn']
         return fqbn
 
-    def _get_flash_settings(self) -> dict[str, str]:
-        """Get flash settings from flash_args file."""
-        flash_args_file = os.path.realpath(os.path.join(self.binary_path, 'flash_args'))
-        with open(flash_args_file) as f:
-            flash_args = f.readline().split(' ')
+    def _parse_flash_args(self) -> tuple[dict[str, str], list[tuple[str, str]]]:
+        """Parse the ``flash_args`` file produced by the Arduino build system.
 
-        flash_settings = {}
-        for i, arg in enumerate(flash_args):
-            if arg.startswith('--'):
-                flash_settings[arg[2:].strip()] = flash_args[i + 1].strip()
+        Returns ``(flash_settings, flash_files)`` where *flash_settings* is a
+        dict of ``--flag value`` pairs (e.g. ``{'flash-mode': 'dio'}``), and
+        *flash_files* is a list of ``(hex_address, absolute_path)`` pairs for
+        each binary that should be flashed.
 
-        if flash_settings == {}:
+        Format of ``flash_args``::
+
+            --flash-mode dio --flash-freq 80m --flash-size 4MB
+            0x0 sketch.ino.bootloader.bin
+            0x8000 sketch.ino.partitions.bin
+            0xe000 boot_app0.bin
+            0x10000 sketch.ino.bin
+        """
+        flash_args_file = os.path.realpath(os.path.join(self.binary_path, 'flash_args'))
+        with open(flash_args_file) as f:
+            lines = f.read().splitlines()
+
+        flash_settings: dict[str, str] = {}
+        flash_files: list[tuple[str, str]] = []
+
+        for line in lines:
+            tokens = line.split()
+            if not tokens:
+                continue
+
+            if tokens[0].startswith('--'):
+                for i, tok in enumerate(tokens):
+                    if tok.startswith('--') and i + 1 < len(tokens):
+                        flash_settings[tok[2:].strip()] = tokens[i + 1].strip()
+            elif _HEX_ADDR_RE.match(tokens[0]) and len(tokens) >= 2:
+                addr = tokens[0]
+                name = tokens[1]
+                path = os.path.realpath(os.path.join(self.binary_path, name))
+                flash_files.append((addr, path))
+
+        if not flash_settings:
             raise ValueError(f'Flash settings not found in {flash_args_file}')
 
-        return flash_settings
+        return flash_settings, flash_files

pytest-embedded-arduino/pytest_embedded_arduino/serial.py

@@ -1,10 +1,20 @@
 import logging
+import shutil
+from pathlib import Path
 
 import esptool
 from pytest_embedded_serial_esp.serial import EspSerial
 
 from .app import ArduinoApp
 
+_ALWAYS_FLASH = {'boot_app0.bin'}
+"""Binaries that must always be fully flashed (never receive a --diff-with ref).
+
+boot_app0.bin selects the active OTA partition slot.  If the user performed an
+OTA update since the last flash, the on-chip copy will differ from the reference
+even though our local copy has not changed.
+"""
+
 
 class ArduinoSerial(EspSerial):
     """
@@ -19,14 +29,26 @@ def __init__(
         self,
         app: ArduinoApp,
         target: str | None = None,
+        fast_flash: bool = True,
         **kwargs,
     ) -> None:
         self.app = app
+        self.fast_flash = fast_flash
         super().__init__(
             target=target or self.app.target,
             **kwargs,
         )
 
+    def _ref_path(self, binary: str) -> Path:
+        """Return the ``*_flashed.bin`` reference path for *binary* inside the build dir."""
+        p = Path(binary)
+        return Path(self.app.binary_path) / (p.stem + '_flashed' + p.suffix)
+
+    @property
+    def _ref_binaries(self) -> list[Path]:
+        """All potential reference files for the current flash_files list."""
+        return [self._ref_path(path) for _, path in self.app.flash_files]
+
     def _start(self):
         if self.skip_autoflash:
             logging.info('Skipping auto flash...')
@@ -37,7 +59,16 @@ def _start(self):
     @EspSerial.use_esptool()
     def flash(self) -> None:
         """
-        Flash the binary files to the board.
+        Flash individual binary files to the board.
+
+        Uses esptool's ``--diff-with`` for fast reflashing when reference binaries
+        from the previous successful flash are available, writing only changed
+        4 KB sectors.  References are saved after each successful flash and
+        invalidated by :meth:`erase_flash`.
+
+        Unlike the merged-binary approach, individual binaries do not overlap
+        with writable flash regions (NVS, OTA data, etc.), so the post-flash
+        MD5 verification succeeds and ``--diff-with`` works correctly.
         """
 
         flash_settings = []
@@ -48,17 +79,78 @@ def flash(self) -> None:
         if self.esp_flash_force:
             flash_settings.append('--force')
 
+        addr_file_pairs: list[str] = []
+        diff_args: list[str] = []
+        have_any_ref = False
+
+        for addr, binary in self.app.flash_files:
+            addr_file_pairs.extend([addr, binary])
+
+            if not self.fast_flash:
+                continue
+
+            name = Path(binary).name
+            if name in _ALWAYS_FLASH:
+                diff_args.append('skip')
+                continue
+
+            ref = self._ref_path(binary)
+            if ref.exists():
+                diff_args.append(str(ref))
+                have_any_ref = True
+            else:
+                diff_args.append('skip')
+
+        if self.fast_flash:
+            if have_any_ref:
+                logging.info('fast-flash: reflashing with references for %d/%d binaries',
+                             sum(1 for d in diff_args if d != 'skip'), len(diff_args))
+            else:
+                diff_args = []
+                logging.info('fast-flash: no references found, performing full flash')
+        else:
+            logging.info('fast-flash: disabled, performing full flash')
+
+        diff_with = ['--diff-with', *diff_args] if diff_args else []
+
         try:
             esptool.main(
                 [
                     '--chip',
                     self.app.target,
                     'write-flash',
-                    '0x0',  # Merged binary is flashed at offset 0
-                    self.app.binary_file,
+                    *addr_file_pairs,
                     *flash_settings,
+                    *diff_with,
                 ],
                 esp=self.esp,
             )
         except Exception:
             raise
+        else:
+            if self.fast_flash:
+                for _, binary in self.app.flash_files:
+                    ref = self._ref_path(binary)
+                    try:
+                        if Path(binary).exists():
+                            shutil.copy2(binary, ref)
+                    except OSError as e:
+                        logging.warning(
+                            'fast-flash: could not save reference for %s (%s)',
+                            Path(binary).name,
+                            e,
+                        )
+
+    def erase_flash(self, force: bool = False) -> None:
+        """
+        Erase the complete flash and invalidate all fast-flash reference binaries.
+        """
+        super().erase_flash(force=force)
+        if self.fast_flash:
+            for ref in self._ref_binaries:
+                if ref.exists():
+                    try:
+                        ref.unlink()
+                        logging.debug('fast-flash: removed reference %s after erase', ref.name)
+                    except OSError as e:
+                        logging.warning('fast-flash: could not remove reference %s (%s)', ref.name, e)