diff --git a/tools/ci/python_packages/ttfw_idf/CIAssignUnitTest.py b/tools/ci/python_packages/ttfw_idf/CIAssignUnitTest.py
index b7f1502e04..bde9a102dc 100644
--- a/tools/ci/python_packages/ttfw_idf/CIAssignUnitTest.py
+++ b/tools/ci/python_packages/ttfw_idf/CIAssignUnitTest.py
@@ -1,7 +1,7 @@
 """
 Command line tool to assign unit tests to CI test jobs.
 """
-
+import os
 import re
 import argparse
 
@@ -145,15 +145,33 @@ class UnitTestAssignTest(CIAssignTest.AssignTest):
         The unit test cases is stored in a yaml file which is created in job build-idf-test.
         """
 
-        try:
-            with open(test_case_path, "r") as f:
-                raw_data = yaml.load(f, Loader=Loader)
-            test_cases = raw_data["test cases"]
-            for case in test_cases:
-                case["tags"] = set(case["tags"])
-        except IOError:
+        def find_by_suffix(suffix, path):
+            res = []
+            for root, _, files in os.walk(path):
+                for file in files:
+                    if file.endswith(suffix):
+                        res.append(os.path.join(root, file))
+            return res
+
+        def get_test_cases_from_yml(yml_file):
+            try:
+                with open(yml_file) as fr:
+                    raw_data = yaml.load(fr, Loader=Loader)
+                test_cases = raw_data['test cases']
+            except (IOError, KeyError):
+                return []
+            else:
+                return test_cases
+
+        test_cases = []
+        if os.path.isdir(test_case_path):
+            for yml_file in find_by_suffix('.yml', test_case_path):
+                test_cases.extend(get_test_cases_from_yml(yml_file))
+        elif os.path.isfile(test_case_path):
+            test_cases.extend(get_test_cases_from_yml(test_case_path))
+        else:
             print("Test case path is invalid. Should only happen when use @bot to skip unit test.")
-            test_cases = []
+
         # filter keys are lower case. Do map lower case keys with original keys.
         try:
             key_mapping = {x.lower(): x for x in test_cases[0].keys()}