From 4558824150a705bd293a232c22ec5494cfac0164 Mon Sep 17 00:00:00 2001
From: "Michael (XIAO Xufeng)" <xiaoxufeng@espressif.com>
Date: Wed, 21 Aug 2019 12:04:26 +0800
Subject: [PATCH] IDFDUT: seperate into different classes

The DUT should be created as the correct sub classes. This can be done
in the config file (UT_xxx_x.yml) Filter
---
 tools/tiny-test-fw/IDF/IDFDUT.py | 39 +++++++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/tools/tiny-test-fw/IDF/IDFDUT.py b/tools/tiny-test-fw/IDF/IDFDUT.py
index c8ee24f23..e44a230d9 100644
--- a/tools/tiny-test-fw/IDF/IDFDUT.py
+++ b/tools/tiny-test-fw/IDF/IDFDUT.py
@@ -122,7 +122,7 @@ def _uses_esptool(func):
         settings = self.port_inst.get_settings()
 
         try:
-            rom = esptool.ESP32ROM(self.port_inst)
+            rom = self._get_rom()(self.port_inst)
             rom.connect('hard_reset')
             esp = rom.run_stub()
 
@@ -159,6 +159,10 @@ class IDFDUT(DUT.SerialDUT):
         self.exceptions = _queue.Queue()
         self.performance_items = _queue.Queue()
 
+    @classmethod
+    def _get_rom(cls):
+        raise NotImplementedError("This is an abstraction class, method not defined.")
+
     @classmethod
     def get_mac(cls, app, port):
         """
@@ -169,7 +173,7 @@ class IDFDUT(DUT.SerialDUT):
         :return: MAC address or None
         """
         try:
-            esp = esptool.ESP32ROM(port)
+            esp = cls._get_rom()(port)
             esp.connect()
             return esp.read_mac()
         except RuntimeError:
@@ -181,7 +185,18 @@ class IDFDUT(DUT.SerialDUT):
 
     @classmethod
     def confirm_dut(cls, port, app, **kwargs):
-        return cls.get_mac(app, port) is not None
+        try:
+            # TODO: check whether 8266 works with this logic
+            # Otherwise overwrite it in ESP8266DUT
+            inst = esptool.ESPLoader.detect_chip(port)
+            if type(inst) != cls._get_rom():
+                raise RuntimeError("Target not expected")
+            return inst.read_mac() is not None
+        except(esptool.FatalError, RuntimeError):
+            return False
+        finally:
+            if inst:
+                inst._port.close()
 
     @_uses_esptool
     def _try_flash(self, esp, erase_nvs, baud_rate):
@@ -389,3 +404,21 @@ class IDFDUT(DUT.SerialDUT):
         if not self.allow_dut_exception and self.get_exceptions():
             Utility.console_log("DUT exception detected on {}".format(self), color="red")
             raise IDFDUTException()
+
+
+class ESP32DUT(IDFDUT):
+    @classmethod
+    def _get_rom(cls):
+        return esptool.ESP32ROM
+
+
+class ESP32S2DUT(IDFDUT):
+    @classmethod
+    def _get_rom(cls):
+        return esptool.ESP32S2ROM
+
+
+class ESP8266DUT(IDFDUT):
+    @classmethod
+    def _get_rom(cls):
+        return esptool.ESP8266ROM