Merge pull request #421 from lucasssvaz/feat/wokwi_jtag

hfudev · web-flow · commit 70f4e74b3b3b · 2026-05-27T01:12:57.000+02:00
feat(wokwi): Add support for using the USB Serial JTAG
diff --git a/pytest-embedded-wokwi/README.md b/pytest-embedded-wokwi/README.md
@@ -30,6 +30,19 @@ To run your tests with Wokwi, make sure to specify the `wokwi` service when runn
 pytest --embedded-services idf,wokwi
 ```
 
+#### USB Serial JTAG
+
+By default, Wokwi diagrams use UART connections (`$serialMonitor:TX`/`$serialMonitor:RX`) for serial communication. Some targets (e.g. ESP32-P4) can use USB Serial JTAG instead. You can enable this with the `--wokwi-usb-serial-jtag` flag:
+
+```
+pytest --embedded-services idf,wokwi --wokwi-usb-serial-jtag true
+```
+
+This works for both auto-generated diagrams and diagrams loaded from disk (including those specified via `--wokwi-diagram`). When enabled, the flag will:
+
+- Set the `serialInterface` attribute to `USB_SERIAL_JTAG` on the board part
+- Remove any `$serialMonitor` connections from the diagram
+
 #### Writing Tests
 
 When writing tests for your firmware, you can use the same pytest fixtures and assertions as you would for local testing. The main difference is that your tests will be executed in the Wokwi simulation environment and you have access to the Wokwi API for controlling the simulation through the `wokwi` fixture.
diff --git a/pytest-embedded-wokwi/pytest_embedded_wokwi/wokwi.py b/pytest-embedded-wokwi/pytest_embedded_wokwi/wokwi.py
@@ -2,6 +2,7 @@
 import logging
 import os
 import sys
+import tempfile
 import typing as t
 from pathlib import Path
 
@@ -47,6 +48,7 @@ def __init__(
         msg_queue: MessageQueue,
         firmware_resolver: IDFFirmwareResolver,
         wokwi_diagram: str | None = None,
+        wokwi_usb_serial_jtag: bool | None = None,
         app: t.Optional['IdfApp'] = None,
         meta: Meta | None = None,
         **kwargs,
@@ -55,6 +57,7 @@ def __init__(
         super().__init__(msg_queue=msg_queue, meta=meta, **kwargs)
 
         self.app = app
+        self._usb_serial_jtag = wokwi_usb_serial_jtag
 
         # Get Wokwi API token
         token = os.getenv('WOKWI_CLI_TOKEN')
@@ -69,6 +72,9 @@ def __init__(
             self.create_diagram_json()
             wokwi_diagram = os.path.join(self.app.app_path, 'diagram.json')
 
+        if self._usb_serial_jtag:
+            wokwi_diagram = self._apply_serial_interface_override(wokwi_diagram)
+
         # Connect and start simulation
         try:
             firmware_path = Path(firmware_resolver.resolve_firmware(app)).as_posix()
@@ -319,6 +325,43 @@ def create_diagram_json(self):
         with open(diagram_json_path, 'w') as f:
             json.dump(diagram, f, indent=2)
 
+    @staticmethod
+    def _apply_serial_interface_override(diagram_path: str) -> str:
+        """Override the serial interface to USB Serial JTAG in a diagram file.
+
+        Reads the diagram JSON, sets the ``serialInterface`` attribute to
+        ``USB_SERIAL_JTAG`` on the main board part and removes any
+        ``$serialMonitor`` connections.  The modified diagram is written to a
+        temporary file so the original is never mutated.
+
+        Returns the path to the modified diagram file.
+        """
+        with open(diagram_path) as f:
+            diagram = json.load(f)
+
+        for part in diagram.get('parts', []):
+            if part.get('type', '').startswith('board-'):
+                part.setdefault('attrs', {})
+                part['attrs']['serialInterface'] = 'USB_SERIAL_JTAG'
+                break
+
+        diagram['connections'] = [
+            conn
+            for conn in diagram.get('connections', [])
+            if not any('$serialMonitor' in str(endpoint) for endpoint in conn)
+        ]
+
+        fd, tmp_path = tempfile.mkstemp(suffix='.json', prefix='wokwi_diagram_')
+        try:
+            with os.fdopen(fd, 'w') as f:
+                json.dump(diagram, f, indent=2)
+        except Exception:
+            os.close(fd)
+            raise
+
+        logging.info('Applied USB Serial JTAG interface override to diagram: %s', tmp_path)
+        return tmp_path
+
     def _hard_reset(self):
         """Fake hard_reset to maintain API consistency."""
         raise NotImplementedError('Hard reset not supported in Wokwi simulation')
diff --git a/pytest-embedded-wokwi/tests/test_wokwi.py b/pytest-embedded-wokwi/tests/test_wokwi.py
@@ -1,6 +1,8 @@
+import json
 import os
 
 import pytest
+from pytest_embedded_wokwi.wokwi import Wokwi
 
 wokwi_token_required = pytest.mark.skipif(
     not os.getenv('WOKWI_CLI_TOKEN', None),
@@ -54,3 +56,115 @@ def test_pexpect_by_wokwi(dut):
     )
 
     result.assert_outcomes(passed=1)
+
+
+class TestApplySerialInterfaceOverride:
+    """Unit tests for Wokwi._apply_serial_interface_override (no token needed)."""
+
+    def _write_diagram(self, tmp_path, diagram: dict) -> str:
+        path = os.path.join(str(tmp_path), 'diagram.json')
+        with open(path, 'w') as f:
+            json.dump(diagram, f)
+        return path
+
+    def test_adds_serial_interface_and_removes_serial_monitor(self, tmp_path):
+        diagram = {
+            'version': 1,
+            'parts': [{'type': 'board-esp32-devkit-c-v4', 'id': 'esp'}],
+            'connections': [
+                ['esp:TX', '$serialMonitor:RX', ''],
+                ['esp:RX', '$serialMonitor:TX', ''],
+            ],
+        }
+        src = self._write_diagram(tmp_path, diagram)
+        result_path = Wokwi._apply_serial_interface_override(src)
+
+        try:
+            with open(result_path) as f:
+                result = json.load(f)
+
+            assert result['parts'][0]['attrs']['serialInterface'] == 'USB_SERIAL_JTAG'
+            assert result['connections'] == []
+        finally:
+            os.unlink(result_path)
+
+    def test_preserves_non_serial_monitor_connections(self, tmp_path):
+        diagram = {
+            'version': 1,
+            'parts': [{'type': 'board-esp32-s3-devkitc-1', 'id': 'esp32', 'attrs': {}}],
+            'connections': [
+                ['esp32:RX', '$serialMonitor:TX', '', []],
+                ['esp32:TX', '$serialMonitor:RX', '', []],
+                ['btn1:1.l', 'esp32:14', 'blue', ['h-38.4', 'v105.78']],
+                ['esp32:4', 'led1:A', 'green', ['h0']],
+            ],
+        }
+        src = self._write_diagram(tmp_path, diagram)
+        result_path = Wokwi._apply_serial_interface_override(src)
+
+        try:
+            with open(result_path) as f:
+                result = json.load(f)
+
+            assert len(result['connections']) == 2
+            assert result['connections'][0] == ['btn1:1.l', 'esp32:14', 'blue', ['h-38.4', 'v105.78']]
+            assert result['connections'][1] == ['esp32:4', 'led1:A', 'green', ['h0']]
+        finally:
+            os.unlink(result_path)
+
+    def test_does_not_mutate_original_file(self, tmp_path):
+        diagram = {
+            'version': 1,
+            'parts': [{'type': 'board-esp32-devkit-c-v4', 'id': 'esp'}],
+            'connections': [
+                ['esp:TX', '$serialMonitor:RX', ''],
+                ['esp:RX', '$serialMonitor:TX', ''],
+            ],
+        }
+        src = self._write_diagram(tmp_path, diagram)
+        result_path = Wokwi._apply_serial_interface_override(src)
+
+        try:
+            with open(src) as f:
+                original = json.load(f)
+
+            assert 'serialInterface' not in original['parts'][0].get('attrs', {})
+            assert len(original['connections']) == 2
+            assert result_path != src
+        finally:
+            os.unlink(result_path)
+
+    def test_adds_attrs_when_missing(self, tmp_path):
+        diagram = {
+            'version': 1,
+            'parts': [{'type': 'board-esp32-p4-function-ev', 'id': 'esp'}],
+            'connections': [],
+        }
+        src = self._write_diagram(tmp_path, diagram)
+        result_path = Wokwi._apply_serial_interface_override(src)
+
+        try:
+            with open(result_path) as f:
+                result = json.load(f)
+
+            assert result['parts'][0]['attrs'] == {'serialInterface': 'USB_SERIAL_JTAG'}
+            assert result['connections'] == []
+        finally:
+            os.unlink(result_path)
+
+    def test_overrides_existing_serial_interface(self, tmp_path):
+        diagram = {
+            'version': 1,
+            'parts': [{'type': 'board-esp32-devkit-c-v4', 'id': 'esp', 'attrs': {'serialInterface': 'UART'}}],
+            'connections': [],
+        }
+        src = self._write_diagram(tmp_path, diagram)
+        result_path = Wokwi._apply_serial_interface_override(src)
+
+        try:
+            with open(result_path) as f:
+                result = json.load(f)
+
+            assert result['parts'][0]['attrs']['serialInterface'] == 'USB_SERIAL_JTAG'
+        finally:
+            os.unlink(result_path)
diff --git a/pytest-embedded/pytest_embedded/dut_factory.py b/pytest-embedded/pytest_embedded/dut_factory.py
@@ -157,6 +157,7 @@ def _fixture_classes_and_options_fn(
     qemu_extra_args,
     qemu_efuse_path,
     wokwi_diagram,
+    wokwi_usb_serial_jtag,
     skip_regenerate_image,
     encrypt,
     keyfile,
@@ -335,6 +336,7 @@ def _fixture_classes_and_options_fn(
                 kwargs[fixture].update(
                     {
                         'wokwi_diagram': wokwi_diagram,
+                        'wokwi_usb_serial_jtag': wokwi_usb_serial_jtag,
                         'msg_queue': msg_queue,
                         'app': None,
                         'meta': _meta,
@@ -696,6 +698,7 @@ def create(
         qemu_extra_args: str | None = None,
         qemu_efuse_path: str | None = None,
         wokwi_diagram: str | None = None,
+        wokwi_usb_serial_jtag: bool | None = None,
         skip_regenerate_image: bool | None = None,
         encrypt: bool | None = None,
         keyfile: str | None = None,
@@ -743,6 +746,7 @@ def create(
             qemu_extra_args: Additional QEMU arguments.
             qemu_efuse_path: Efuse binary path.
             wokwi_diagram: Wokwi diagram path.
+            wokwi_usb_serial_jtag: Use USB Serial JTAG instead of UART for Wokwi serial communication.
             skip_regenerate_image: Skip image regeneration flag.
             encrypt: Encryption flag.
             keyfile: Keyfile for encryption.
@@ -813,6 +817,7 @@ def create(
                 'qemu_extra_args': qemu_extra_args,
                 'qemu_efuse_path': qemu_efuse_path,
                 'wokwi_diagram': wokwi_diagram,
+                'wokwi_usb_serial_jtag': wokwi_usb_serial_jtag,
                 'skip_regenerate_image': skip_regenerate_image,
                 'encrypt': encrypt,
                 'keyfile': keyfile,
diff --git a/pytest-embedded/pytest_embedded/plugin.py b/pytest-embedded/pytest_embedded/plugin.py
@@ -298,6 +298,13 @@ def pytest_addoption(parser):
         '--wokwi-diagram',
         help='Path to the wokwi diagram file (Default: None)',
     )
+    wokwi_group.addoption(
+        '--wokwi-usb-serial-jtag',
+        help='y/yes/true for True and n/no/false for False. '
+        'Use USB Serial JTAG instead of UART for serial communication in the Wokwi diagram. '
+        'When enabled, the diagram will use the USB_SERIAL_JTAG interface and remove $serialMonitor connections. '
+        '(Default: False)',
+    )
 
 
 ###########
@@ -1076,6 +1083,13 @@ def wokwi_diagram(request: FixtureRequest) -> str | None:
     return _request_param_or_config_option_or_default(request, 'wokwi_diagram', None)
 
 
+@pytest.fixture
+@multi_dut_argument
+def wokwi_usb_serial_jtag(request: FixtureRequest) -> bool | None:
+    """Enable parametrization for the same cli option"""
+    return _request_param_or_config_option_or_default(request, 'wokwi_usb_serial_jtag', None)
+
+
 ####################
 # Private Fixtures #
 ####################
@@ -1134,6 +1148,7 @@ def parametrize_fixtures(
     qemu_extra_args,
     qemu_efuse_path,
     wokwi_diagram,
+    wokwi_usb_serial_jtag,
     skip_regenerate_image,
     encrypt,
     keyfile,
