dict directly by using a set in discover_unit_tests

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -305,7 +305,7 @@ def discover_unit_tests(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] | None = None,
-) -> tuple[dict[str, list[FunctionCalledInTest]], int]:
+) -> tuple[dict[str, set[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:
@@ -326,7 +326,7 @@ def discover_tests_pytest(
     cfg: TestConfig,
     discover_only_these_tests: list[Path] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> dict[str, set[FunctionCalledInTest]]:
     tests_root = cfg.tests_root
     project_root = cfg.project_root_path
 
@@ -402,7 +402,7 @@ def discover_tests_unittest(
     cfg: TestConfig,
     discover_only_these_tests: list[str] | None = None,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> dict[str, set[FunctionCalledInTest]]:
     tests_root: Path = cfg.tests_root
     loader: unittest.TestLoader = unittest.TestLoader()
     tests: unittest.TestSuite = loader.discover(str(tests_root))
@@ -469,7 +469,7 @@ def process_test_files(
     file_to_test_map: dict[Path, list[TestsInFile]],
     cfg: TestConfig,
     functions_to_optimize: list[FunctionToOptimize] | None = None,
-) -> dict[str, list[FunctionCalledInTest]]:
+) -> dict[str, set[FunctionCalledInTest]]:
     import jedi
 
     project_root_path = cfg.project_root_path
@@ -637,4 +637,4 @@ def process_test_files(
 
             progress.advance(task_id)
 
-    return {function: list(tests) for function, tests in function_to_test_map.items()}
+    return dict(function_to_test_map)

codeflash/optimization/function_optimizer.py

@@ -57,7 +57,6 @@
 from codeflash.models.models import (
     BestOptimization,
     CodeOptimizationContext,
-    FunctionCalledInTest,
     GeneratedTests,
     GeneratedTestsList,
     OptimizationSet,
@@ -87,7 +86,13 @@
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.either import Result
-    from codeflash.models.models import BenchmarkKey, CoverageData, FunctionSource, OptimizedCandidate
+    from codeflash.models.models import (
+        BenchmarkKey,
+        CoverageData,
+        FunctionCalledInTest,
+        FunctionSource,
+        OptimizedCandidate,
+    )
     from codeflash.verification.verification_utils import TestConfig
 
 
@@ -97,7 +102,7 @@ def __init__(
         function_to_optimize: FunctionToOptimize,
         test_cfg: TestConfig,
         function_to_optimize_source_code: str = "",
-        function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None,
+        function_to_tests: dict[str, set[FunctionCalledInTest]] | None = None,
         function_to_optimize_ast: ast.FunctionDef | None = None,
         aiservice_client: AiServiceClient | None = None,
         function_benchmark_timings: dict[BenchmarkKey, int] | None = None,
@@ -213,7 +218,7 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
 
         function_to_optimize_qualified_name = self.function_to_optimize.qualified_name
         function_to_all_tests = {
-            key: self.function_to_tests.get(key, []) + function_to_concolic_tests.get(key, [])
+            key: self.function_to_tests.get(key, set()) | function_to_concolic_tests.get(key, set())
             for key in set(self.function_to_tests) | set(function_to_concolic_tests)
         }
         instrumented_unittests_created_for_function = self.instrument_existing_tests(function_to_all_tests)
@@ -690,7 +695,7 @@ def cleanup_leftover_test_return_values() -> None:
         get_run_tmp_file(Path("test_return_values_0.bin")).unlink(missing_ok=True)
         get_run_tmp_file(Path("test_return_values_0.sqlite")).unlink(missing_ok=True)
 
-    def instrument_existing_tests(self, function_to_all_tests: dict[str, list[FunctionCalledInTest]]) -> set[Path]:
+    def instrument_existing_tests(self, function_to_all_tests: dict[str, set[FunctionCalledInTest]]) -> set[Path]:
         existing_test_files_count = 0
         replay_test_files_count = 0
         concolic_coverage_test_files_count = 0
@@ -701,7 +706,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi
             logger.info(f"Did not find any pre-existing tests for '{func_qualname}', will only use generated tests.")
             console.rule()
         else:
-            test_file_invocation_positions = defaultdict(list[FunctionCalledInTest])
+            test_file_invocation_positions = defaultdict(list)
             for tests_in_file in function_to_all_tests.get(func_qualname):
                 test_file_invocation_positions[
                     (tests_in_file.tests_in_file.test_file, tests_in_file.tests_in_file.test_type)
@@ -787,7 +792,7 @@ def generate_tests_and_optimizations(
         generated_test_paths: list[Path],
         generated_perf_test_paths: list[Path],
         run_experiment: bool = False,  # noqa: FBT001, FBT002
-    ) -> Result[tuple[GeneratedTestsList, dict[str, list[FunctionCalledInTest]], OptimizationSet], str]:
+    ) -> Result[tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet], str]:
         assert len(generated_test_paths) == N_TESTS_TO_GENERATE
         max_workers = N_TESTS_TO_GENERATE + 2 if not run_experiment else N_TESTS_TO_GENERATE + 3
         console.rule()

codeflash/optimization/optimizer.py

@@ -48,7 +48,7 @@ def create_function_optimizer(
         self,
         function_to_optimize: FunctionToOptimize,
         function_to_optimize_ast: ast.FunctionDef | None = None,
-        function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None,
+        function_to_tests: dict[str, set[FunctionCalledInTest]] | None = None,
         function_to_optimize_source_code: str | None = "",
         function_benchmark_timings: dict[str, dict[BenchmarkKey, float]] | None = None,
         total_benchmark_timings: dict[BenchmarkKey, float] | None = None,

codeflash/result/create_pr.py

@@ -25,7 +25,7 @@
 
 def existing_tests_source_for(
     function_qualified_name_with_modules_from_root: str,
-    function_to_tests: dict[str, list[FunctionCalledInTest]],
+    function_to_tests: dict[str, set[FunctionCalledInTest]],
     tests_root: Path,
 ) -> str:
     test_files = function_to_tests.get(function_qualified_name_with_modules_from_root)

codeflash/verification/concolic_testing.py

@@ -24,7 +24,7 @@
 
 def generate_concolic_tests(
     test_cfg: TestConfig, args: Namespace, function_to_optimize: FunctionToOptimize, function_to_optimize_ast: ast.AST
-) -> tuple[dict[str, list[FunctionCalledInTest]], str]:
+) -> tuple[dict[str, set[FunctionCalledInTest]], str]:
     start_time = time.perf_counter()
     function_to_concolic_tests = {}
     concolic_test_suite_code = ""

tests/test_unit_test_discovery.py

@@ -129,12 +129,9 @@ def test_discover_tests_pytest_with_temp_dir_root():
         # Check if the dummy test file is discovered
         assert len(discovered_tests) == 1
         assert len(discovered_tests["dummy_code.dummy_function"]) == 2
-        assert discovered_tests["dummy_code.dummy_function"][0].tests_in_file.test_file == test_file_path
-        assert discovered_tests["dummy_code.dummy_function"][1].tests_in_file.test_file == test_file_path
-        assert {
-            discovered_tests["dummy_code.dummy_function"][0].tests_in_file.test_function,
-            discovered_tests["dummy_code.dummy_function"][1].tests_in_file.test_function,
-        } == {"test_dummy_parametrized_function[True]", "test_dummy_function"}
+        dummy_tests = discovered_tests["dummy_code.dummy_function"]
+        assert all(test.tests_in_file.test_file == test_file_path for test in dummy_tests)
+        assert {test.tests_in_file.test_function for test in dummy_tests} == {"test_dummy_parametrized_function[True]", "test_dummy_function"}
 
 
 def test_discover_tests_pytest_with_multi_level_dirs():
@@ -201,13 +198,13 @@ def test_discover_tests_pytest_with_multi_level_dirs():
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 3
-        assert discovered_tests["root_code.root_function"][0].tests_in_file.test_file == root_test_file_path
+        assert next(iter(discovered_tests["root_code.root_function"])).tests_in_file.test_file == root_test_file_path
         assert (
-            discovered_tests["level1.level1_code.level1_function"][0].tests_in_file.test_file == level1_test_file_path
+            next(iter(discovered_tests["level1.level1_code.level1_function"])).tests_in_file.test_file == level1_test_file_path
         )
 
         assert (
-            discovered_tests["level1.level2.level2_code.level2_function"][0].tests_in_file.test_file
+            next(iter(discovered_tests["level1.level2.level2_code.level2_function"])).tests_in_file.test_file
             == level2_test_file_path
         )
 
@@ -291,17 +288,17 @@ def test_discover_tests_pytest_dirs():
 
         # Check if the test files at all levels are discovered
         assert len(discovered_tests) == 4
-        assert discovered_tests["root_code.root_function"][0].tests_in_file.test_file == root_test_file_path
+        assert next(iter(discovered_tests["root_code.root_function"])).tests_in_file.test_file == root_test_file_path
         assert (
-            discovered_tests["level1.level1_code.level1_function"][0].tests_in_file.test_file == level1_test_file_path
+            next(iter(discovered_tests["level1.level1_code.level1_function"])).tests_in_file.test_file == level1_test_file_path
         )
         assert (
-            discovered_tests["level1.level2.level2_code.level2_function"][0].tests_in_file.test_file
+            next(iter(discovered_tests["level1.level2.level2_code.level2_function"])).tests_in_file.test_file
             == level2_test_file_path
         )
 
         assert (
-            discovered_tests["level1.level3.level3_code.level3_function"][0].tests_in_file.test_file
+            next(iter(discovered_tests["level1.level3.level3_code.level3_function"])).tests_in_file.test_file
             == level3_test_file_path
         )
 
@@ -337,7 +334,7 @@ def test_discover_tests_pytest_with_class():
 
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
-        assert discovered_tests["some_class_code.SomeClass.some_method"][0].tests_in_file.test_file == test_file_path
+        assert next(iter(discovered_tests["some_class_code.SomeClass.some_method"])).tests_in_file.test_file == test_file_path
 
 
 def test_discover_tests_pytest_with_double_nested_directories():
@@ -376,9 +373,7 @@ def test_discover_tests_pytest_with_double_nested_directories():
         # Check if the test class and method are discovered
         assert len(discovered_tests) == 1
         assert (
-            discovered_tests["nested.more_nested.nested_class_code.NestedClass.nested_method"][
-                0
-            ].tests_in_file.test_file
+            next(iter(discovered_tests["nested.more_nested.nested_class_code.NestedClass.nested_method"])).tests_in_file.test_file
             == test_file_path
         )
 
@@ -425,7 +420,7 @@ def test_discover_tests_with_code_in_dir_and_test_in_subdir():
 
         # Check if the test file is discovered and associated with the code file
         assert len(discovered_tests) == 1
-        assert discovered_tests["code.some_code.some_function"][0].tests_in_file.test_file == test_file_path
+        assert next(iter(discovered_tests["code.some_code.some_function"])).tests_in_file.test_file == test_file_path
 
 
 def test_discover_tests_pytest_with_nested_class():
@@ -465,7 +460,7 @@ def test_discover_tests_pytest_with_nested_class():
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
         assert (
-            discovered_tests["nested_class_code.OuterClass.InnerClass.inner_method"][0].tests_in_file.test_file
+            next(iter(discovered_tests["nested_class_code.OuterClass.InnerClass.inner_method"])).tests_in_file.test_file
             == test_file_path
         )
 
@@ -504,7 +499,7 @@ def test_discover_tests_pytest_separate_moduledir():
 
         # Check if the test for the nested class method is discovered
         assert len(discovered_tests) == 1
-        assert discovered_tests["mypackage.code.find_common_tags"][0].tests_in_file.test_file == test_file_path
+        assert next(iter(discovered_tests["mypackage.code.find_common_tags"])).tests_in_file.test_file == test_file_path
 
 
 def test_unittest_discovery_with_pytest():
@@ -548,8 +543,9 @@ def test_add(self):
         assert len(discovered_tests) == 1
         assert "calculator.Calculator.add" in discovered_tests
         assert len(discovered_tests["calculator.Calculator.add"]) == 1
-        assert discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_file == test_file_path
-        assert discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_function == "test_add"
+        calculator_test = next(iter(discovered_tests["calculator.Calculator.add"]))
+        assert calculator_test.tests_in_file.test_file == test_file_path
+        assert calculator_test.tests_in_file.test_function == "test_add"
 
 
 def test_unittest_discovery_with_pytest_parent_class():
@@ -615,8 +611,9 @@ def test_add(self):
         assert len(discovered_tests) == 2
         assert "calculator.Calculator.add" in discovered_tests
         assert len(discovered_tests["calculator.Calculator.add"]) == 1
-        assert discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_file == test_file_path
-        assert discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_function == "test_add"
+        calculator_test = next(iter(discovered_tests["calculator.Calculator.add"]))
+        assert calculator_test.tests_in_file.test_file == test_file_path
+        assert calculator_test.tests_in_file.test_function == "test_add"
 
 
 def test_unittest_discovery_with_pytest_private():
@@ -712,9 +709,10 @@ def test_add_with_parameters(self):
         assert len(discovered_tests) == 1
         assert "calculator.Calculator.add" in discovered_tests
         assert len(discovered_tests["calculator.Calculator.add"]) == 1
-        assert discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_file == test_file_path
+        calculator_test = next(iter(discovered_tests["calculator.Calculator.add"]))
+        assert calculator_test.tests_in_file.test_file == test_file_path
         assert (
-            discovered_tests["calculator.Calculator.add"][0].tests_in_file.test_function == "test_add_with_parameters"
+            calculator_test.tests_in_file.test_function == "test_add_with_parameters"
         )