From f137ce0213fc6b396593d3b0664483bff4c2537d Mon Sep 17 00:00:00 2001
From: horw <horw7001@gmail.com>
Date: Wed, 7 Feb 2024 15:45:14 +0800
Subject: [PATCH] feat: add esp-flash-force

---
 .../pytest_embedded_arduino/serial.py         |  3 ++
 .../pytest_embedded_idf/serial.py             |  5 ++-
 pytest-embedded-idf/tests/test_idf.py         | 37 +++++++++++++++++++
 pytest-embedded/pytest_embedded/plugin.py     | 16 +++++++-
 4 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/pytest-embedded-arduino/pytest_embedded_arduino/serial.py b/pytest-embedded-arduino/pytest_embedded_arduino/serial.py
index a7f71989..1b80de17 100644
--- a/pytest-embedded-arduino/pytest_embedded_arduino/serial.py
+++ b/pytest-embedded-arduino/pytest_embedded_arduino/serial.py
@@ -51,6 +51,9 @@ def flash(self) -> None:
             flash_settings.append(f'--{k}')
             flash_settings.append(v)
 
+        if self.app.esp_flash_force:
+            flash_settings.append('--force')
+
         try:
             esptool.main(
                 ['--chip', self.app.target, 'write_flash', *flash_files, *flash_settings],
diff --git a/pytest-embedded-idf/pytest_embedded_idf/serial.py b/pytest-embedded-idf/pytest_embedded_idf/serial.py
index e1f59396..02ddc28e 100644
--- a/pytest-embedded-idf/pytest_embedded_idf/serial.py
+++ b/pytest-embedded-idf/pytest_embedded_idf/serial.py
@@ -109,7 +109,10 @@ def load_ram(self) -> None:
             )
 
     def _force_flag(self):
+        if self.app.esp_flash_force:
+            return ['--force']
         config = self.app.sdkconfig
+        print('CONFIG', config)
         if any((config.get('CONFIG_SECURE_FLASH_ENC_ENABLED', False), config.get('CONFIG_SECURE_BOOT', False))):
             return ['--force']
         return []
@@ -178,7 +181,7 @@ def flash(self) -> None:
 
         _args.extend(self.app.flash_args['write_flash_args'])
         _args.extend(self._force_flag())
-
+        print('ARGS', _args)
         esptool.main(_args, esp=self.esp)
 
         if self._meta:
diff --git a/pytest-embedded-idf/tests/test_idf.py b/pytest-embedded-idf/tests/test_idf.py
index 8c01292f..a3dd9da8 100644
--- a/pytest-embedded-idf/tests/test_idf.py
+++ b/pytest-embedded-idf/tests/test_idf.py
@@ -35,6 +35,43 @@ def test_idf_serial_flash(dut):
     result.assert_outcomes(passed=1)
 
 
+def test_esp_flash_force_flag(testdir):
+    testdir.makepyfile("""
+        import pexpect
+        import pytest
+
+        def test_idf_serial_flash(dut):
+            dut.expect('Hello world!')
+            assert dut.app.esp_flash_force == True
+    """)
+    result = testdir.runpytest(
+        '-s',
+        '--embedded-services', 'esp,idf',
+        '--app-path', os.path.join(testdir.tmpdir, 'hello_world_esp32'),
+        '--esp-flash-force',
+    )
+
+    result.assert_outcomes(passed=1)
+
+
+def test_esp_flash_no_force_flag(testdir):
+    testdir.makepyfile("""
+        import pexpect
+        import pytest
+
+        def test_idf_serial_flash(dut):
+            dut.expect('Hello world!')
+            assert dut.app.esp_flash_force == False
+    """)
+    result = testdir.runpytest(
+        '-s',
+        '--embedded-services', 'esp,idf',
+        '--app-path', os.path.join(testdir.tmpdir, 'hello_world_esp32'),
+    )
+
+    result.assert_outcomes(passed=1)
+
+
 def test_expect_no_matching(testdir):
     testdir.makepyfile("""
         import pexpect
diff --git a/pytest-embedded/pytest_embedded/plugin.py b/pytest-embedded/pytest_embedded/plugin.py
index eb5e42b8..f0ac04f9 100644
--- a/pytest-embedded/pytest_embedded/plugin.py
+++ b/pytest-embedded/pytest_embedded/plugin.py
@@ -164,7 +164,11 @@ def pytest_addoption(parser):
         '--port-mac',
         help='MAC address of the board. (Default: None)',
     )
-
+    esp_group.addoption(
+        '--esp-flash-force',
+        action='store_true',
+        help='force mode for esptool',
+    )
     idf_group = parser.getgroup('embedded-idf')
     idf_group.addoption(
         '--part-tool',
@@ -765,6 +769,13 @@ def app_path(request: FixtureRequest, test_file_path: str) -> t.Optional[str]:
     return _request_param_or_config_option_or_default(request, 'app_path', os.path.dirname(test_file_path))
 
 
+@pytest.fixture
+@multi_dut_argument
+def esp_flash_force(request: FixtureRequest) -> t.Optional[str]:
+    """Enable parametrization for the same cli option"""
+    return _request_param_or_config_option_or_default(request, 'esp_flash_force', False)
+
+
 @pytest.fixture
 @multi_dut_argument
 def build_dir(request: FixtureRequest) -> t.Optional[str]:
@@ -1024,6 +1035,7 @@ def _fixture_classes_and_options(
     skip_autoflash,
     erase_all,
     esptool_baud,
+    esp_flash_force,
     part_tool,
     confirm_target_elf_sha256,
     erase_nvs,
@@ -1074,7 +1086,7 @@ def _fixture_classes_and_options(
 
     for fixture in FIXTURES_SERVICES.keys():
         if fixture == 'app':
-            kwargs['app'] = {'app_path': app_path, 'build_dir': build_dir}
+            kwargs['app'] = {'app_path': app_path, 'build_dir': build_dir, 'esp_flash_force': esp_flash_force}
             if 'idf' in _services:
                 if 'qemu' in _services:
                     from pytest_embedded_qemu import DEFAULT_IMAGE_FN, QemuApp