unit tests: Fix retrying with a new baud rate in IDFDUT.start_app()

diff --git a/tools/tiny-test-fw/IDF/IDFDUT.py b/tools/tiny-test-fw/IDF/IDFDUT.py
index 418b9c1671d90da222b11200c225b25801a79d2f..b4a8d5674af894da4a720606f6062fc9c26df41c 100644 (file)
--- a/tools/tiny-test-fw/IDF/IDFDUT.py
+++ b/tools/tiny-test-fw/IDF/IDFDUT.py
@@ -104,60 +104,69 @@ class IDFDUT(DUT.SerialDUT):
         return cls.get_mac(app, port) is not None
 
     @_uses_esptool
-    def start_app(self, esp, erase_nvs=ERASE_NVS):
+    def _try_flash(self, esp, erase_nvs, baud_rate):
         """
-        download and start app.
+        Called by start_app() to try flashing at a particular baud rate.
 
-        :param: erase_nvs: whether erase NVS partition during flash
-        :return: None
+        Structured this way so @_uses_esptool will reconnect each time
         """
-        flash_files = [ (offs, open(path, "rb")) for (offs, path) in self.app.flash_files ]
-
-        if erase_nvs:
-            address = self.app.partition_table["nvs"]["offset"]
-            size = self.app.partition_table["nvs"]["size"]
-            nvs_file = tempfile.TemporaryFile()
-            nvs_file.write(b'\xff' * size)
-            nvs_file.seek(0)
-            flash_files.append( (int(address, 0), nvs_file) )
-
-        # fake flasher args object, this is a hack until
-        # esptool Python API is improved
-        Flash_Args = namedtuple('write_flash_args',
-                                ['flash_size',
-                                 'flash_mode',
-                                 'flash_freq',
-                                 'addr_filename',
-                                 'no_stub',
-                                 'compress',
-                                 'verify',
-                                 'encrypt'])
-
-        flash_args = Flash_Args(
-            self.app.flash_settings["flash_size"],
-            self.app.flash_settings["flash_mode"],
-            self.app.flash_settings["flash_freq"],
-            flash_files,
-            False,
-            True,
-            False,
-            False
-        )
-
         try:
-            for baud_rate in [ 921600, 115200 ]:
-                try:
-                    esp.change_baud(baud_rate)
-                    esptool.write_flash(esp, flash_args)
-                    break
-                except RuntimeError:
-                    continue
-            else:
-                raise IDFToolError()
+            # note: opening here prevents us from having to seek back to 0 each time
+            flash_files = [ (offs, open(path, "rb")) for (offs, path) in self.app.flash_files ]
+
+            if erase_nvs:
+                address = self.app.partition_table["nvs"]["offset"]
+                size = self.app.partition_table["nvs"]["size"]
+                nvs_file = tempfile.TemporaryFile()
+                nvs_file.write(b'\xff' * size)
+                nvs_file.seek(0)
+                flash_files.append( (int(address, 0), nvs_file) )
+
+            # fake flasher args object, this is a hack until
+            # esptool Python API is improved
+            Flash_Args = namedtuple('write_flash_args',
+                                    ['flash_size',
+                                     'flash_mode',
+                                     'flash_freq',
+                                     'addr_filename',
+                                     'no_stub',
+                                     'compress',
+                                     'verify',
+                                     'encrypt'])
+
+            flash_args = Flash_Args(
+                self.app.flash_settings["flash_size"],
+                self.app.flash_settings["flash_mode"],
+                self.app.flash_settings["flash_freq"],
+                flash_files,
+                False,
+                True,
+                False,
+                False
+            )
+
+            esp.change_baud(baud_rate)
+            esptool.write_flash(esp, flash_args)
         finally:
             for (_,f) in flash_files:
                 f.close()
 
+    def start_app(self, erase_nvs=ERASE_NVS):
+        """
+        download and start app.
+
+        :param: erase_nvs: whether erase NVS partition during flash
+        :return: None
+        """
+        for baud_rate in [ 921600, 115200 ]:
+            try:
+                self._try_flash(erase_nvs, baud_rate)
+                break
+            except RuntimeError:
+                continue
+        else:
+            raise IDFToolError()
+
     @_uses_esptool
     def reset(self, esp):
         """