diff --git a/examples/storage/parttool/example_test.py b/examples/storage/parttool/example_test.py
index 17f8d01af..6ba90c175 100644
--- a/examples/storage/parttool/example_test.py
+++ b/examples/storage/parttool/example_test.py
@@ -29,9 +29,9 @@ def test_examples_parttool(env, extra_data):
     script_path = os.path.join(os.getenv("IDF_PATH"), "examples", "storage", "parttool", "parttool_example.py")
 
     binary_path = ""
-    for config in dut.download_config:
-        if "parttool.bin" in config:
-            binary_path = config
+    for flash_file in dut.app.flash_files:
+        if "parttool.bin" in flash_file[1]:
+            binary_path = flash_file[1]
             break
 
     subprocess.check_call([sys.executable, script_path, "--binary", binary_path])
diff --git a/examples/system/ota/otatool/example_test.py b/examples/system/ota/otatool/example_test.py
index 0cfa44ef1..a22ebf5a8 100644
--- a/examples/system/ota/otatool/example_test.py
+++ b/examples/system/ota/otatool/example_test.py
@@ -32,9 +32,9 @@ def test_otatool_example(env, extra_data):
     script_path = os.path.join(os.getenv("IDF_PATH"), "examples", "system", "ota", "otatool", "otatool_example.py")
     binary_path = ""
 
-    for config in dut.download_config:
-        if "otatool.bin" in config:
-            binary_path = config
+    for flash_file in dut.app.flash_files:
+        if "otatool.bin" in flash_file[1]:
+            binary_path = flash_file[1]
             break
 
     subprocess.check_call([sys.executable, script_path, "--binary", binary_path])
diff --git a/tools/tiny-test-fw/IDF/IDFDUT.py b/tools/tiny-test-fw/IDF/IDFDUT.py
index 9d2f25caf..ec690a177 100644
--- a/tools/tiny-test-fw/IDF/IDFDUT.py
+++ b/tools/tiny-test-fw/IDF/IDFDUT.py
@@ -22,8 +22,6 @@ import tempfile
 
 from serial.tools import list_ports
 
-from collections import namedtuple
-
 import DUT
 
 try:
@@ -51,14 +49,20 @@ def _uses_esptool(func):
 
         settings = self.port_inst.get_settings()
 
-        rom = esptool.ESP32ROM(self.port_inst)
-        rom.connect('hard_reset')
-        esp = rom.run_stub()
+        try:
+            rom = esptool.ESP32ROM(self.port_inst)
+            rom.connect('hard_reset')
+            esp = rom.run_stub()
 
-        ret = func(self, esp, *args, **kwargs)
+            ret = func(self, esp, *args, **kwargs)
+            # do hard reset after use esptool
+            esp.hard_reset()
+        finally:
+            # always need to restore port settings
+            self.port_inst.apply_settings(settings)
 
-        self.port_inst.apply_settings(settings)
         self.start_receive()
+
         return ret
     return handler
 
@@ -94,6 +98,8 @@ class IDFDUT(DUT.SerialDUT):
         except RuntimeError:
             return None
         finally:
+            # do hard reset after use esptool
+            esp.hard_reset()
             esp._port.close()
 
     @classmethod
@@ -121,31 +127,27 @@ class IDFDUT(DUT.SerialDUT):
 
             # 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'])
+            class FlashArgs(object):
+                def __init__(self, attributes):
+                    for key, value in attributes.items():
+                        self.__setattr__(key, value)
 
-            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
-            )
+            flash_args = FlashArgs({
+                'flash_size': self.app.flash_settings["flash_size"],
+                'flash_mode': self.app.flash_settings["flash_mode"],
+                'flash_freq': self.app.flash_settings["flash_freq"],
+                'addr_filename': flash_files,
+                'no_stub': False,
+                'compress': True,
+                'verify': False,
+                'encrypt': False,
+            })
 
             esp.change_baud(baud_rate)
+            esptool.detect_flash_size(esp, flash_args)
             esptool.write_flash(esp, flash_args)
         finally:
-            for (_,f) in flash_files:
+            for (_, f) in flash_files:
                 f.close()
 
     def start_app(self, erase_nvs=ERASE_NVS):
@@ -171,7 +173,9 @@ class IDFDUT(DUT.SerialDUT):
 
         :return: None
         """
-        esp.hard_reset()
+        # decorator `_use_esptool` will do reset
+        # so we don't need to do anything in this method
+        pass
 
     @_uses_esptool
     def erase_partition(self, esp, partition):