From a5bcaea9cc69f888198f8af2a9552557e6e23356 Mon Sep 17 00:00:00 2001
From: Fu Hanxi <fuhanxi@espressif.com>
Date: Thu, 3 Dec 2020 16:32:54 +0800
Subject: [PATCH] use "additional_duts" to overwrite dut class for different
 targets

---
 tools/ci/python_packages/ttfw_idf/__init__.py | 34 +++++++++++--------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/tools/ci/python_packages/ttfw_idf/__init__.py b/tools/ci/python_packages/ttfw_idf/__init__.py
index d033cbd878..03673807b6 100644
--- a/tools/ci/python_packages/ttfw_idf/__init__.py
+++ b/tools/ci/python_packages/ttfw_idf/__init__.py
@@ -16,6 +16,7 @@ import json
 import logging
 import os
 import re
+from copy import deepcopy
 
 import junit_xml
 
@@ -73,11 +74,11 @@ def local_test_check(decorator_target):
         try:
             idf_target = sdkconfig['IDF_TARGET'].upper()
         except KeyError:
-            logging.warning('IDF_TARGET not in {}. IDF_TARGET set to esp32'.format(os.path.abspath(expected_json_path)))
+            logging.debug('IDF_TARGET not in {}. IDF_TARGET set to esp32'.format(os.path.abspath(expected_json_path)))
         else:
-            logging.info('IDF_TARGET: {}'.format(idf_target))
+            logging.debug('IDF_TARGET: {}'.format(idf_target))
     else:
-        logging.warning('{} not found. IDF_TARGET set to esp32'.format(os.path.abspath(expected_json_path)))
+        logging.debug('{} not found. IDF_TARGET set to esp32'.format(os.path.abspath(expected_json_path)))
 
     if isinstance(decorator_target, list):
         if idf_target not in decorator_target:
@@ -88,11 +89,11 @@ def local_test_check(decorator_target):
     return idf_target
 
 
-def get_dut_class(target, erase_nvs=None):
-    if target not in TARGET_DUT_CLS_DICT:
-        raise Exception('target can only be {%s} (case insensitive)' % ', '.join(TARGET_DUT_CLS_DICT.keys()))
+def get_dut_class(target, dut_class_dict, erase_nvs=None):
+    if target not in dut_class_dict:
+        raise Exception('target can only be {%s} (case insensitive)' % ', '.join(dut_class_dict.keys()))
 
-    dut = TARGET_DUT_CLS_DICT[target.upper()]
+    dut = dut_class_dict[target.upper()]
     if erase_nvs:
         dut.ERASE_NVS = 'erase_nvs'
     return dut
@@ -111,16 +112,19 @@ def ci_target_check(func):
     return wrapper
 
 
-def test_func_generator(func, app, target, ci_target, module, execution_time, level, erase_nvs, drop_kwargs_dut=False, **kwargs):
+def test_func_generator(func, app, target, ci_target, module, execution_time, level, erase_nvs, **kwargs):
+    target = upper_list_or_str(target)
     test_target = local_test_check(target)
-    dut = get_dut_class(test_target, erase_nvs)
-    if drop_kwargs_dut and 'dut' in kwargs:  # panic_test() will inject dut, resolve conflicts here
-        dut = kwargs['dut']
-        del kwargs['dut']
+    if 'additional_duts' in kwargs:
+        dut_classes = deepcopy(TARGET_DUT_CLS_DICT)
+        dut_classes.update(kwargs['additional_duts'])
+    else:
+        dut_classes = TARGET_DUT_CLS_DICT
+    dut = get_dut_class(test_target, dut_classes, erase_nvs)
     original_method = TinyFW.test_method(
-        app=app, dut=dut, target=upper_list_or_str(target), ci_target=upper_list_or_str(ci_target),
+        app=app, dut=dut, target=target, ci_target=upper_list_or_str(ci_target),
         module=module, execution_time=execution_time, level=level, erase_nvs=erase_nvs,
-        dut_dict=TARGET_DUT_CLS_DICT, **kwargs
+        dut_dict=dut_classes, **kwargs
     )
     test_func = original_method(func)
     test_func.case_info["ID"] = format_case_id(target, test_func.case_info["name"])
@@ -188,7 +192,7 @@ def idf_custom_test(app=TestApp, target="ESP32", ci_target=None, module="misc",
     :return: test method
     """
     def test(func):
-        return test_func_generator(func, app, target, ci_target, module, execution_time, level, erase_nvs, drop_kwargs_dut=True, **kwargs)
+        return test_func_generator(func, app, target, ci_target, module, execution_time, level, erase_nvs, **kwargs)
     return test