go

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -139,7 +139,7 @@ def close(self) -> None:
 
 def discover_unit_tests(
     cfg: TestConfig, discover_only_these_tests: list[Path] | None = None
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
     framework_strategies: dict[str, Callable] = {"pytest": discover_tests_pytest, "unittest": discover_tests_unittest}
     strategy = framework_strategies.get(cfg.test_framework, None)
     if not strategy:
@@ -151,7 +151,7 @@ def discover_unit_tests(
 
 def discover_tests_pytest(
     cfg: TestConfig, discover_only_these_tests: list[Path] | None = None
-) -> dict[Path, list[FunctionCalledInTest]]:
+) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
 
@@ -225,7 +225,7 @@ def discover_tests_pytest(
 
 def discover_tests_unittest(
     cfg: TestConfig, discover_only_these_tests: list[str] | None = None
-) -> dict[Path, list[FunctionCalledInTest]]:
+) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
     tests: unittest.TestSuite = loader.discover(str(tests_root))
@@ -290,15 +290,14 @@ def discover_parameters_unittest(function_name: str) -> tuple[bool, str, str | N
 
 def process_test_files(
     file_to_test_map: dict[Path, list[TestsInFile]], cfg: TestConfig
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
     import jedi
 
     project_root_path = cfg.project_root_path
     test_framework = cfg.test_framework
 
     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 (
@@ -412,13 +411,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
@@ -476,4 +470,6 @@ def process_test_files(
             progress.advance(task_id)
 
     tests_cache.close()
-    return {function: list(tests) for function, tests in function_to_test_map.items()}
+    function_to_tests_dict = {function: list(tests) for function, tests in function_to_test_map.items()}
+    num_discovered_tests = sum(len(tests) for tests in function_to_tests_dict.values())
+    return function_to_tests_dict, num_discovered_tests

codeflash/discovery/functions_to_optimize.py

@@ -268,7 +268,7 @@ def find_all_functions_in_file(file_path: Path) -> dict[Path, list[FunctionToOpt
 def get_all_replay_test_functions(
     replay_test: Path, test_cfg: TestConfig, project_root_path: Path
 ) -> dict[Path, list[FunctionToOptimize]]:
-    function_tests = discover_unit_tests(test_cfg, discover_only_these_tests=[replay_test])
+    function_tests, _ = discover_unit_tests(test_cfg, discover_only_these_tests=[replay_test])
     # Get the absolute file paths for each function, excluding class name if present
     filtered_valid_functions = defaultdict(list)
     file_to_functions_map = defaultdict(list)

codeflash/optimization/optimizer.py

@@ -162,8 +162,8 @@ def run(self) -> None:
 
             console.rule()
             start_time = time.time()
-            function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg)
-            num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()])
+            function_to_tests: dict[str, list[FunctionCalledInTest]]
+            function_to_tests, num_discovered_tests = discover_unit_tests(self.test_cfg)
             console.rule()
             logger.info(
                 f"Discovered {num_discovered_tests} existing unit tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"

codeflash/verification/concolic_testing.py

@@ -78,8 +78,7 @@ def generate_concolic_tests(
                 test_framework=args.test_framework,
                 pytest_cmd=args.pytest_cmd,
             )
-            function_to_concolic_tests = discover_unit_tests(concolic_test_cfg)
-            num_discovered_concolic_tests: int = sum([len(value) for value in function_to_concolic_tests.values()])
+            function_to_concolic_tests, num_discovered_concolic_tests = discover_unit_tests(concolic_test_cfg)
             logger.info(
                 f"Created {num_discovered_concolic_tests} "
                 f"concolic unit test case{'s' if num_discovered_concolic_tests != 1 else ''} "

tests/test_unit_test_discovery.py

@@ -15,7 +15,7 @@ def test_unit_test_discovery_pytest():
         test_framework="pytest",
         tests_project_rootdir=tests_path.parent,
     )
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     assert len(tests) > 0
 
 
@@ -28,7 +28,7 @@ def test_benchmark_test_discovery_pytest():
         test_framework="pytest",
         tests_project_rootdir=tests_path.parent,
     )
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     assert len(tests) == 1 # Should not discover benchmark tests
 
 
@@ -42,7 +42,7 @@ def test_unit_test_discovery_unittest():
         tests_project_rootdir=project_path.parent,
     )
     os.chdir(project_path)
-    tests = discover_unit_tests(test_config)
+    tests, _ = discover_unit_tests(test_config)
     # assert len(tests) > 0
     # Unittest discovery within a pytest environment does not work
 
@@ -80,7 +80,7 @@ def sorter(arr):
         )
 
         # Discover tests
-        tests = discover_unit_tests(test_config)
+        tests, _ = discover_unit_tests(test_config)
         assert len(tests) == 1
         assert 'bubble_sort.sorter' in tests
         assert len(tests['bubble_sort.sorter']) == 2
@@ -119,7 +119,7 @@ def test_discover_tests_pytest_with_temp_dir_root():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the dummy test file is discovered
         assert len(discovered_tests) == 1
@@ -192,7 +192,7 @@ def test_discover_tests_pytest_with_multi_level_dirs():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 3
@@ -282,7 +282,7 @@ def test_discover_tests_pytest_dirs():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 4
@@ -328,7 +328,7 @@ def test_discover_tests_pytest_with_class():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
@@ -366,7 +366,7 @@ def test_discover_tests_pytest_with_double_nested_directories():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
@@ -416,7 +416,7 @@ def test_discover_tests_with_code_in_dir_and_test_in_subdir():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test file is discovered and associated with the code file
         assert len(discovered_tests) == 1
@@ -455,7 +455,7 @@ def test_discover_tests_pytest_with_nested_class():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
@@ -495,7 +495,7 @@ def test_discover_tests_pytest_separate_moduledir():
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
@@ -537,7 +537,7 @@ def test_add(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 1
@@ -604,7 +604,7 @@ def test_add(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 2
@@ -649,7 +649,7 @@ def _test_add(self):  # Private test method should not be discovered
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify no tests were discovered
         assert len(discovered_tests) == 0
@@ -701,7 +701,7 @@ def test_add_with_parameters(self):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the unittest was discovered
         assert len(discovered_tests) == 1
@@ -783,7 +783,7 @@ def test_add_mixed(self, name, a, b, expected):
         )
 
         # Discover tests
-        discovered_tests = discover_unit_tests(test_config)
+        discovered_tests, _ = discover_unit_tests(test_config)
 
         # Verify the basic structure
         assert len(discovered_tests) == 2  # Should have tests for both add and multiply