From 9a9fc1d61a95784f46d29214a5796d8624eee98d Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Thu, 20 Nov 2025 17:29:45 -0800 Subject: [PATCH 1/6] testing with fixtures --- tests/test_unit_test_discovery.py | 85 ++++++++++++++++++++++++++++++- 1 file changed, 84 insertions(+), 1 deletion(-) diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 1a0c73ffa..dd5883503 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -8,7 +8,7 @@ filter_test_files_by_imports, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize -from codeflash.models.models import TestsInFile, TestType +from codeflash.models.models import TestsInFile, TestType, FunctionParent from codeflash.verification.verification_utils import TestConfig @@ -714,6 +714,61 @@ def test_add_with_parameters(self): assert calculator_test.tests_in_file.test_file.resolve() == test_file_path.resolve() assert calculator_test.tests_in_file.test_function == "test_add_with_parameters" +def test_unittest_discovery_with_pytest_fixture(): + with tempfile.TemporaryDirectory() as tmpdirname: + path_obj_tmpdirname = Path(tmpdirname) + + # Create a simple code file + code_file_path = path_obj_tmpdirname / "topological_sort.py" + code_file_content = """ +import uuid +from collections import defaultdict + + +class Graph: + def __init__(self, vertices: int): + self.vertices=vertices + + def topologicalSort(self): + return self.vertices + +""" + code_file_path.write_text(code_file_content) + + # Create a unittest test file with parameterized tests + test_file_path = path_obj_tmpdirname / "test_topological_sort.py" + test_file_content = """ +from topological_sort import Graph +import pytest + +@pytest.fixture +def g(self): + return Graph(6) + +def test_topological_sort(g): + assert g.topologicalSort() == 6 +""" + test_file_path.write_text(test_file_content) + + # Configure test discovery + test_config = TestConfig( + tests_root=path_obj_tmpdirname, + project_root_path=path_obj_tmpdirname, + test_framework="pytest", # Using pytest framework to discover unittest tests + tests_project_rootdir=path_obj_tmpdirname.parent, + ) + fto = FunctionToOptimize(function_name="topologicalSort", file_path=code_file_path, parents=[FunctionParent(name="Graph", type="ClassDef")]) + # Discover tests + discovered_tests, _, _ = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file_path: [fto]}) + + # Verify the unittest was discovered + assert len(discovered_tests) == 1 + assert "topological_sort.Graph.topologicalSort" in discovered_tests + assert len(discovered_tests["topological_sort.Graph.topologicalSort"]) == 1 + tpsort_test = next(iter(discovered_tests["topological_sort.Graph.topologicalSort"])) + assert tpsort_test.tests_in_file.test_file.resolve() == test_file_path.resolve() + assert tpsort_test.tests_in_file.test_function == "test_topological_sort" + def test_unittest_discovery_with_pytest_parameterized(): with tempfile.TemporaryDirectory() as tmpdirname: @@ -1335,6 +1390,34 @@ def test_topological_sort(): assert should_process is True +def test_analyze_imports_fixture(): + with tempfile.TemporaryDirectory() as tmpdirname: + test_file = Path(tmpdirname) / "test_example.py" + test_content = """ +from code_to_optimize.topological_sort import Graph +import pytest + +@pytest.fixture +def g(self): + return Graph(6) + +def test_topological_sort(g): + g.addEdge(5, 2) + g.addEdge(5, 0) + g.addEdge(4, 0) + g.addEdge(4, 1) + g.addEdge(2, 3) + g.addEdge(3, 1) + + assert g.topologicalSort()[0] == [5, 4, 2, 3, 1, 0] +""" + test_file.write_text(test_content) + + target_functions = {"Graph.topologicalSort"} + should_process = analyze_imports_in_test_file(test_file, target_functions) + + assert should_process is True + def test_analyze_imports_aliased_class_method_negative(): with tempfile.TemporaryDirectory() as tmpdirname: test_file = Path(tmpdirname) / "test_example.py" From 9ac79fb1f39be57a92d6ffe168b26efdd6e43f01 Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Thu, 20 Nov 2025 17:48:59 -0800 Subject: [PATCH 2/6] only capturing for dummy_fn but not tpsort --- tests/test_unit_test_discovery.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index dd5883503..b4bfee9e7 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -729,6 +729,9 @@ class Graph: def __init__(self, vertices: int): self.vertices=vertices + def dummy_fn(self): + return 1 + def topologicalSort(self): return self.vertices @@ -746,6 +749,7 @@ def g(self): return Graph(6) def test_topological_sort(g): + assert g.dummy_fn() == 1 assert g.topologicalSort() == 6 """ test_file_path.write_text(test_file_content) From 105735fd4f512bd87d93dfd2949d4149bc3b9564 Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Thu, 20 Nov 2025 18:48:05 -0800 Subject: [PATCH 3/6] only capturing for dummy_fn but not tpsort --- tests/test_unit_test_discovery.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index b4bfee9e7..5805eefd0 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -745,7 +745,7 @@ def topologicalSort(self): import pytest @pytest.fixture -def g(self): +def g(): return Graph(6) def test_topological_sort(g): @@ -766,7 +766,7 @@ def test_topological_sort(g): discovered_tests, _, _ = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file_path: [fto]}) # Verify the unittest was discovered - assert len(discovered_tests) == 1 + assert len(discovered_tests) == 2 assert "topological_sort.Graph.topologicalSort" in discovered_tests assert len(discovered_tests["topological_sort.Graph.topologicalSort"]) == 1 tpsort_test = next(iter(discovered_tests["topological_sort.Graph.topologicalSort"])) @@ -1402,7 +1402,7 @@ def test_analyze_imports_fixture(): import pytest @pytest.fixture -def g(self): +def g(): return Graph(6) def test_topological_sort(g): From 86a8d202a640f2f01e72fb961c3d3396fc2dd9f1 Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Thu, 20 Nov 2025 19:03:17 -0800 Subject: [PATCH 4/6] possible fix --- codeflash/discovery/discover_unit_tests.py | 52 ++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index d4bc9abd7..bc0e2fd67 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -868,6 +868,58 @@ def process_test_files( continue try: if not definition or definition[0].type != "function": + # Fallback: Try to match against functions_to_optimize when Jedi can't resolve + # This handles cases where Jedi fails with pytest fixtures + if functions_to_optimize and name.name: + for func_to_opt in functions_to_optimize: + # Check if this unresolved name matches a function we're looking for + if func_to_opt.function_name == name.name: + # Check if the test file imports the class/module containing this function + qualified_name_with_modules = func_to_opt.qualified_name_with_modules_from_root( + project_root_path + ) + + # Only add if this test actually tests the function we're optimizing + for test_func in test_functions_by_name[scope]: + if test_func.parameters is not None: + if test_framework == "pytest": + scope_test_function = ( + f"{test_func.function_name}[{test_func.parameters}]" + ) + else: # unittest + scope_test_function = ( + f"{test_func.function_name}_{test_func.parameters}" + ) + else: + scope_test_function = test_func.function_name + + function_to_test_map[qualified_name_with_modules].add( + FunctionCalledInTest( + tests_in_file=TestsInFile( + test_file=test_file, + test_class=test_func.test_class, + test_function=scope_test_function, + test_type=test_func.test_type, + ), + position=CodePosition(line_no=name.line, col_no=name.column), + ) + ) + tests_cache.insert_test( + file_path=str(test_file), + file_hash=file_hash, + qualified_name_with_modules_from_root=qualified_name_with_modules, + function_name=scope, + test_class=test_func.test_class or "", + test_function=scope_test_function, + test_type=test_func.test_type, + line_number=name.line, + col_number=name.column, + ) + + if test_func.test_type == TestType.REPLAY_TEST: + num_discovered_replay_tests += 1 + + num_discovered_tests += 1 continue definition_obj = definition[0] definition_path = str(definition_obj.module_path) From a6f7d5a953b1198583d5af42d6cf48ded72dbd2b Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Fri, 21 Nov 2025 11:25:49 -0800 Subject: [PATCH 5/6] more tests --- tests/test_unit_test_discovery.py | 240 ++++++++++++++++++++++++++++++ 1 file changed, 240 insertions(+) diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 5805eefd0..4534b6ce0 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -773,6 +773,203 @@ def test_topological_sort(g): assert tpsort_test.tests_in_file.test_file.resolve() == test_file_path.resolve() assert tpsort_test.tests_in_file.test_function == "test_topological_sort" +def test_unittest_discovery_with_pytest_class_fixture(): + with tempfile.TemporaryDirectory() as tmpdirname: + path_obj_tmpdirname = Path(tmpdirname) + + # Create a simple code file + code_file_path = path_obj_tmpdirname / "router_file.py" + code_file_content = """ +from __future__ import annotations + +import hashlib +import json +from typing import Any, Dict, List, Optional, Required, TypedDict, Union # noqa: UP035 + + +class LiteLLMParamsTypedDict(TypedDict, total=False): + model: str + custom_llm_provider: Optional[str] + tpm: Optional[int] + rpm: Optional[int] + order: Optional[int] + weight: Optional[int] + max_parallel_requests: Optional[int] + api_key: Optional[str] + api_base: Optional[str] + api_version: Optional[str] + stream_timeout: Optional[Union[float, str]] + max_retries: Optional[int] + organization: Optional[Union[list, str]] # for openai orgs + ## DROP PARAMS ## + drop_params: Optional[bool] + ## UNIFIED PROJECT/REGION ## + region_name: Optional[str] + ## VERTEX AI ## + vertex_project: Optional[str] + vertex_location: Optional[str] + ## AWS BEDROCK / SAGEMAKER ## + aws_access_key_id: Optional[str] + aws_secret_access_key: Optional[str] + aws_region_name: Optional[str] + ## IBM WATSONX ## + watsonx_region_name: Optional[str] + ## CUSTOM PRICING ## + input_cost_per_token: Optional[float] + output_cost_per_token: Optional[float] + input_cost_per_second: Optional[float] + output_cost_per_second: Optional[float] + num_retries: Optional[int] + ## MOCK RESPONSES ## + + # routing params + # use this for tag-based routing + tags: Optional[list[str]] + + # deployment budgets + max_budget: Optional[float] + budget_duration: Optional[str] + +class DeploymentTypedDict(TypedDict, total=False): + model_name: Required[str] + litellm_params: Required[LiteLLMParamsTypedDict] + model_info: dict + +class Router: + model_names: set = set() # noqa: RUF012 + cache_responses: Optional[bool] = False + default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour + tenacity = None + + def __init__( # noqa: PLR0915 + self, + model_list: Optional[ + Union[list[DeploymentTypedDict], list[dict[str, Any]]] + ] = None, + ) -> None: + self.model_list = model_list # noqa: ARG002 + self.model_id_to_deployment_index_map: dict[str, int] = {} + self.model_name_to_deployment_indices: dict[str, list[int]] = {} + def _generate_model_id(self, model_group: str, litellm_params: dict): # noqa: ANN202 + # Optimized: Use list and join instead of string concatenation in loop + # This avoids creating many temporary string objects (O(n) vs O(n²) complexity) + parts = [model_group] + for k, v in litellm_params.items(): + if isinstance(k, str): + parts.append(k) + elif isinstance(k, dict): + parts.append(json.dumps(k)) + else: + parts.append(str(k)) + + if isinstance(v, str): + parts.append(v) + elif isinstance(v, dict): + parts.append(json.dumps(v)) + else: + parts.append(str(v)) + + concat_str = "".join(parts) + hash_object = hashlib.sha256(concat_str.encode()) + + return hash_object.hexdigest() + def _add_model_to_list_and_index_map( + self, model: dict, model_id: Optional[str] = None + ) -> None: + idx = len(self.model_list) + self.model_list.append(model) + + # Update model_id index for O(1) lookup + if model_id is not None: + self.model_id_to_deployment_index_map[model_id] = idx + elif model.get("model_info", {}).get("id") is not None: + self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx + + # Update model_name index for O(1) lookup + model_name = model.get("model_name") + if model_name: + if model_name not in self.model_name_to_deployment_indices: + self.model_name_to_deployment_indices[model_name] = [] + self.model_name_to_deployment_indices[model_name].append(idx) + + def _build_model_id_to_deployment_index_map(self, model_list: list) -> None: + # First populate the model_list + self.model_list = [] + for _, model in enumerate(model_list): + # Extract model_info from the model dict + model_info = model.get("model_info", {}) + model_id = model_info.get("id") + + # If no ID exists, generate one using the same logic as set_model_list + if model_id is None: + model_name = model.get("model_name", "") + litellm_params = model.get("litellm_params", {}) + model_id = self._generate_model_id(model_name, litellm_params) + # Update the model_info in the original list + if "model_info" not in model: + model["model_info"] = {} + model["model_info"]["id"] = model_id + + self._add_model_to_list_and_index_map(model=model, model_id=model_id) +""" + code_file_path.write_text(code_file_content) + + # Create a unittest test file with parameterized tests + test_file_path = path_obj_tmpdirname / "test_router_file.py" + test_file_content = """ +import pytest + +from router_file import Router + + +class TestRouterIndexManagement: + @pytest.fixture + def router(self): # noqa: ANN201 + return Router(model_list=[]) + def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001 + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "model-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "model-2"}, + }, + ] + + # Test: Build index from model list + router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001 + + # Verify: model_list is populated + assert len(router.model_list) == 2 + # Verify: model_id_to_deployment_index_map is correctly built + assert router.model_id_to_deployment_index_map["model-1"] == 0 + assert router.model_id_to_deployment_index_map["model-2"] == 1 +""" + test_file_path.write_text(test_file_content) + + # Configure test discovery + test_config = TestConfig( + tests_root=path_obj_tmpdirname, + project_root_path=path_obj_tmpdirname, + test_framework="pytest", # Using pytest framework to discover unittest tests + tests_project_rootdir=path_obj_tmpdirname.parent, + ) + fto = FunctionToOptimize(function_name="_build_model_id_to_deployment_index_map", file_path=code_file_path, parents=[FunctionParent(name="Router", type="ClassDef")]) + # Discover tests + discovered_tests, _, _ = discover_unit_tests(test_config, file_to_funcs_to_optimize={code_file_path: [fto]}) + + # Verify the unittest was discovered + assert len(discovered_tests) == 1 + assert "router_file.Router._build_model_id_to_deployment_index_map" in discovered_tests + assert len(discovered_tests["router_file.Router._build_model_id_to_deployment_index_map"]) == 1 + router_test = next(iter(discovered_tests["router_file.Router._build_model_id_to_deployment_index_map"])) + assert router_test.tests_in_file.test_file.resolve() == test_file_path.resolve() + assert router_test.tests_in_file.test_function == "test_build_model_id_to_deployment_index_map" + def test_unittest_discovery_with_pytest_parameterized(): with tempfile.TemporaryDirectory() as tmpdirname: @@ -1422,6 +1619,49 @@ def test_topological_sort(g): assert should_process is True +def test_analyze_imports_class_fixture(): + with tempfile.TemporaryDirectory() as tmpdirname: + test_file = Path(tmpdirname) / "test_example.py" + test_content = """ +import pytest + +from router_file import Router + + +class TestRouterIndexManagement: + @pytest.fixture + def router(self): # noqa: ANN201 + return Router(model_list=[]) + def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001 + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "model-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "model-2"}, + }, + ] + + # Test: Build index from model list + router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001 + + # Verify: model_list is populated + assert len(router.model_list) == 2 + # Verify: model_id_to_deployment_index_map is correctly built + assert router.model_id_to_deployment_index_map["model-1"] == 0 + assert router.model_id_to_deployment_index_map["model-2"] == 1 +""" + test_file.write_text(test_content) + + target_functions = {"Router._build_model_id_to_deployment_index_map"} + should_process = analyze_imports_in_test_file(test_file, target_functions) + + assert should_process is True + def test_analyze_imports_aliased_class_method_negative(): with tempfile.TemporaryDirectory() as tmpdirname: test_file = Path(tmpdirname) / "test_example.py" From 69ef7cc1cae55efff258f875642816c73b8a319f Mon Sep 17 00:00:00 2001 From: Codeflash Bot Date: Fri, 21 Nov 2025 11:51:06 -0800 Subject: [PATCH 6/6] type hints not well supported in 3.9 --- tests/test_unit_test_discovery.py | 84 ++++++------------------------- 1 file changed, 16 insertions(+), 68 deletions(-) diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 4534b6ce0..60facb17b 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -784,73 +784,20 @@ def test_unittest_discovery_with_pytest_class_fixture(): import hashlib import json -from typing import Any, Dict, List, Optional, Required, TypedDict, Union # noqa: UP035 - - -class LiteLLMParamsTypedDict(TypedDict, total=False): - model: str - custom_llm_provider: Optional[str] - tpm: Optional[int] - rpm: Optional[int] - order: Optional[int] - weight: Optional[int] - max_parallel_requests: Optional[int] - api_key: Optional[str] - api_base: Optional[str] - api_version: Optional[str] - stream_timeout: Optional[Union[float, str]] - max_retries: Optional[int] - organization: Optional[Union[list, str]] # for openai orgs - ## DROP PARAMS ## - drop_params: Optional[bool] - ## UNIFIED PROJECT/REGION ## - region_name: Optional[str] - ## VERTEX AI ## - vertex_project: Optional[str] - vertex_location: Optional[str] - ## AWS BEDROCK / SAGEMAKER ## - aws_access_key_id: Optional[str] - aws_secret_access_key: Optional[str] - aws_region_name: Optional[str] - ## IBM WATSONX ## - watsonx_region_name: Optional[str] - ## CUSTOM PRICING ## - input_cost_per_token: Optional[float] - output_cost_per_token: Optional[float] - input_cost_per_second: Optional[float] - output_cost_per_second: Optional[float] - num_retries: Optional[int] - ## MOCK RESPONSES ## - - # routing params - # use this for tag-based routing - tags: Optional[list[str]] - - # deployment budgets - max_budget: Optional[float] - budget_duration: Optional[str] - -class DeploymentTypedDict(TypedDict, total=False): - model_name: Required[str] - litellm_params: Required[LiteLLMParamsTypedDict] - model_info: dict class Router: - model_names: set = set() # noqa: RUF012 - cache_responses: Optional[bool] = False - default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour + model_names: list + cache_responses = False tenacity = None def __init__( # noqa: PLR0915 self, - model_list: Optional[ - Union[list[DeploymentTypedDict], list[dict[str, Any]]] - ] = None, + model_list = None, ) -> None: - self.model_list = model_list # noqa: ARG002 - self.model_id_to_deployment_index_map: dict[str, int] = {} - self.model_name_to_deployment_indices: dict[str, list[int]] = {} - def _generate_model_id(self, model_group: str, litellm_params: dict): # noqa: ANN202 + self.model_list = model_list + self.model_id_to_deployment_index_map = {} + self.model_name_to_deployment_indices = {} + def _generate_model_id(self, model_group, litellm_params): # Optimized: Use list and join instead of string concatenation in loop # This avoids creating many temporary string objects (O(n) vs O(n²) complexity) parts = [model_group] @@ -874,7 +821,7 @@ def _generate_model_id(self, model_group: str, litellm_params: dict): # noqa: A return hash_object.hexdigest() def _add_model_to_list_and_index_map( - self, model: dict, model_id: Optional[str] = None + self, model, model_id = None ) -> None: idx = len(self.model_list) self.model_list.append(model) @@ -892,7 +839,7 @@ def _add_model_to_list_and_index_map( self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) - def _build_model_id_to_deployment_index_map(self, model_list: list) -> None: + def _build_model_id_to_deployment_index_map(self, model_list): # First populate the model_list self.model_list = [] for _, model in enumerate(model_list): @@ -911,6 +858,7 @@ def _build_model_id_to_deployment_index_map(self, model_list: list) -> None: model["model_info"]["id"] = model_id self._add_model_to_list_and_index_map(model=model, model_id=model_id) + """ code_file_path.write_text(code_file_content) @@ -924,9 +872,9 @@ def _build_model_id_to_deployment_index_map(self, model_list: list) -> None: class TestRouterIndexManagement: @pytest.fixture - def router(self): # noqa: ANN201 + def router(self): return Router(model_list=[]) - def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001 + def test_build_model_id_to_deployment_index_map(self, router): model_list = [ { "model_name": "gpt-3.5-turbo", @@ -941,7 +889,7 @@ def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ] # Test: Build index from model list - router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001 + router._build_model_id_to_deployment_index_map(model_list) # Verify: model_list is populated assert len(router.model_list) == 2 @@ -1630,9 +1578,9 @@ def test_analyze_imports_class_fixture(): class TestRouterIndexManagement: @pytest.fixture - def router(self): # noqa: ANN201 + def router(self): return Router(model_list=[]) - def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ANN001 + def test_build_model_id_to_deployment_index_map(self, router): model_list = [ { "model_name": "gpt-3.5-turbo", @@ -1647,7 +1595,7 @@ def test_build_model_id_to_deployment_index_map(self, router) -> None: # noqa: ] # Test: Build index from model list - router._build_model_id_to_deployment_index_map(model_list) # noqa: SLF001 + router._build_model_id_to_deployment_index_map(model_list) # Verify: model_list is populated assert len(router.model_list) == 2