implement suggestions

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -207,11 +207,10 @@ def process_test_files(
     with test_files_progress_bar(
         total=len(file_to_test_map), description="Processing test files"
     ) as (progress, task_id):
-
         for test_file, functions in file_to_test_map.items():
             try:
                 script = jedi.Script(path=test_file, project=jedi_project)
-                test_functions = set()
+                test_functions_by_name: dict[str, TestFunction] = {}
 
                 all_names = script.get_names(all_scopes=True, references=True)
                 all_defs = script.get_names(all_scopes=True, definitions=True)
@@ -238,22 +237,20 @@ def process_test_files(
                             function.test_function
                         )[1]
                         if function_name in top_level_functions:
-                            test_functions.add(
-                                TestFunction(
-                                    function_name,
-                                    function.test_class,
-                                    parameters,
-                                    function.test_type,
-                                )
-                            )
-                    elif function.test_function in top_level_functions:
-                        test_functions.add(
-                            TestFunction(
-                                function.test_function,
+                            qualified_name = function_name
+                            test_functions_by_name[qualified_name] = TestFunction(
+                                function_name,
                                 function.test_class,
-                                None,
+                                parameters,
                                 function.test_type,
                             )
+                    elif function.test_function in top_level_functions:
+                        qualified_name = function.test_function
+                        test_functions_by_name[qualified_name] = TestFunction(
+                            function.test_function,
+                            function.test_class,
+                            None,
+                            function.test_type,
                         )
                     elif UNITTEST_PARAMETERIZED_TEST_NAME_REGEX.match(
                         function.test_function
@@ -262,13 +259,12 @@ def process_test_files(
                             "", function.test_function
                         )
                         if base_name in top_level_functions:
-                            test_functions.add(
-                                TestFunction(
-                                    function_name=base_name,
-                                    test_class=function.test_class,
-                                    parameters=function.test_function,
-                                    test_type=function.test_type,
-                                )
+                            qualified_name = base_name
+                            test_functions_by_name[qualified_name] = TestFunction(
+                                function_name=base_name,
+                                test_class=function.test_class,
+                                parameters=function.test_function,
+                                test_type=function.test_type,
                             )
 
             elif test_framework == "unittest":
@@ -289,7 +285,8 @@ def process_test_files(
                                 )
 
                                 if is_parameterized and new_function == def_name.name:
-                                    test_functions.add(
+                                    qualified_name = f"{matched_name}.{def_name.name}"
+                                    test_functions_by_name[qualified_name] = (
                                         TestFunction(
                                             function_name=def_name.name,
                                             test_class=matched_name,
@@ -298,7 +295,8 @@ def process_test_files(
                                         )
                                     )
                                 elif function == def_name.name:
-                                    test_functions.add(
+                                    qualified_name = f"{matched_name}.{def_name.name}"
+                                    test_functions_by_name[qualified_name] = (
                                         TestFunction(
                                             function_name=def_name.name,
                                             test_class=matched_name,
@@ -307,13 +305,6 @@ def process_test_files(
                                         )
                                     )
 
-            test_functions_list = list(test_functions)
-            test_functions_raw = [elem.function_name for elem in test_functions_list]
-
-            test_functions_by_name = defaultdict(list)
-            for i, func_name in enumerate(test_functions_raw):
-                test_functions_by_name[func_name].append(i)
-
             for name in all_names:
                 if name.full_name is None:
                     continue
@@ -347,22 +338,25 @@ def process_test_files(
                     and definition[0].module_name != name.module_name
                     and definition[0].full_name is not None
                 ):
-                    for index in test_functions_by_name[scope]:
-                        scope_test_function = test_functions_list[index].function_name
-                        scope_test_class = test_functions_list[index].test_class
-                        scope_parameters = test_functions_list[index].parameters
-                        test_type = test_functions_list[index].test_type
-
-                        if scope_parameters is not None:
-                            if test_framework == "pytest":
-                                scope_test_function += "[" + scope_parameters + "]"
-                            if test_framework == "unittest":
-                                scope_test_function += "_" + scope_parameters
+                    test_function_obj = test_functions_by_name[scope]
+                    scope_test_function = test_function_obj.function_name
+                    scope_test_class = test_function_obj.test_class
+                    scope_parameters = test_function_obj.parameters
+                    test_type = test_function_obj.test_type
+
+                    if scope_parameters is not None:
+                        if test_framework == "pytest":
+                            scope_test_function += "[" + scope_parameters + "]"
 
                         full_name_without_module_prefix = definition[
                             0
                         ].full_name.replace(definition[0].module_name + ".", "", 1)
-                        qualified_name_with_modules_from_root = f"{module_name_from_file_path(definition[0].module_path, project_root_path)}.{full_name_without_module_prefix}"
+                        module_name = module_name_from_file_path(
+                            definition[0].module_path, project_root_path
+                        )
+                        qualified_name_with_modules_from_root = (
+                            f"{module_name}.{full_name_without_module_prefix}"
+                        )
 
                         function_to_test_map[qualified_name_with_modules_from_root].add(
                             FunctionCalledInTest(