fix tests & remove goto cache

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -315,10 +315,8 @@ def discover_unit_tests(
     # Extract all functions to optimize for import filtering
     functions_to_optimize = None
     if file_to_funcs_to_optimize:
-        functions_to_optimize = [
-            func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list
-        ]
-    
+        functions_to_optimize = [func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list]
+
     return strategy(cfg, discover_only_these_tests, functions_to_optimize)
 
 
@@ -493,7 +491,6 @@ def process_test_files(
 
     function_to_test_map = defaultdict(set)
     jedi_project = jedi.Project(path=project_root_path)
-    goto_cache = {}
     tests_cache = TestsCache()
 
     with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
@@ -607,13 +604,8 @@ def process_test_files(
                 if scope not in test_functions_by_name:
                     continue
 
-                cache_key = (name.full_name, name.module_name)
                 try:
-                    if cache_key in goto_cache:
-                        definition = goto_cache[cache_key]
-                    else:
-                        definition = name.goto(follow_imports=True, follow_builtin_imports=False)
-                        goto_cache[cache_key] = definition
+                    definition = name.goto(follow_imports=True, follow_builtin_imports=False)
                 except Exception as e:
                     logger.debug(str(e))
                     continue

tests/test_unit_test_discovery.py

@@ -1079,13 +1079,9 @@ def test_other():
         mock_function.function_name = "target_function"
         mock_function.parents = []  # No parent classes
 
-        filtered_tests = discover_unit_tests(test_config, functions_to_optimize=[mock_function])
-        # The import filter is designed for high recall, so it may include both functions
-        # because both test files import from the same module (mycode) that contains target_function
-        assert len(filtered_tests) >= 1  # Should find at least target_function
+        filtered_tests = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file: [mock_function]})
+        assert len(filtered_tests) >= 1
         assert "mycode.target_function" in filtered_tests
-        # In a perfect world we'd filter out other_function, but conservative filtering
-        # is acceptable for performance optimization purposes
 
 
 def test_analyze_imports_conditional_import():
@@ -1225,8 +1221,8 @@ def test_unrelated():
         mock_function.qualified_name_with_modules_from_root.return_value = "target_module.target_function"
         mock_function.function_name = "target_function"
         mock_function.parents = []  # No parent classes
-        
-        filtered_tests = discover_unit_tests(test_config, functions_to_optimize=[mock_function])
+
+        filtered_tests = discover_unit_tests(test_config, file_to_funcs_to_optimize={target_file: [mock_function]})
         # Should filter out the unrelated test since it imports from a different module
         assert len(filtered_tests) == 1
         assert "target_module.target_function" in filtered_tests