Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions codeflash/discovery/discover_unit_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,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, set[FunctionCalledInTest]], int]:
) -> tuple[dict[str, set[FunctionCalledInTest]], int, 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:
Expand All @@ -352,8 +352,10 @@ def discover_unit_tests(
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]
function_to_tests, num_discovered_tests = strategy(cfg, discover_only_these_tests, functions_to_optimize)
return function_to_tests, num_discovered_tests
function_to_tests, num_discovered_tests, num_discovered_replay_tests = strategy(
cfg, discover_only_these_tests, functions_to_optimize
)
return function_to_tests, num_discovered_tests, num_discovered_replay_tests


def discover_tests_pytest(
Expand Down Expand Up @@ -515,6 +517,7 @@ def process_test_files(

function_to_test_map = defaultdict(set)
num_discovered_tests = 0
num_discovered_replay_tests = 0
jedi_project = jedi.Project(path=project_root_path)

with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
Expand Down Expand Up @@ -661,11 +664,14 @@ def process_test_files(
position=CodePosition(line_no=name.line, col_no=name.column),
)
)
if test_func.test_type == TestType.REPLAY_TEST:
num_discovered_replay_tests += 1

num_discovered_tests += 1
except Exception as e:
logger.debug(str(e))
continue

progress.advance(task_id)

return dict(function_to_test_map), num_discovered_tests
return dict(function_to_test_map), num_discovered_tests, num_discovered_replay_tests
2 changes: 1 addition & 1 deletion codeflash/discovery/functions_to_optimize.py
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,7 @@ def get_all_replay_test_functions(
logger.error("Could not find trace_file_path in replay test files.")
exit_with_message("Could not find trace_file_path in replay test files.")

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)
Expand Down
4 changes: 2 additions & 2 deletions codeflash/optimization/optimizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,12 +226,12 @@ def discover_tests(

console.rule()
start_time = time.time()
function_to_tests, num_discovered_tests = discover_unit_tests(
function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
)
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}"
f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
)
console.rule()
ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
Expand Down
Loading
Loading