diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py index 35150e0da..a5be1f01a 100644 --- a/codeflash/code_utils/edit_generated_tests.py +++ b/codeflash/code_utils/edit_generated_tests.py @@ -5,7 +5,7 @@ import re from pathlib import Path from textwrap import dedent -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING import libcst as cst @@ -50,230 +50,178 @@ class CfoVisitor(ast.NodeVisitor): and reports their location relative to the function they're in. """ - def __init__(self, source_code: str) -> None: + def __init__(self, function_name: str, source_code: str) -> None: self.source_lines = source_code.splitlines() + self.name = function_name self.results: list[int] = [] # map actual line number to line number in ast - def _is_codeflash_output_target(self, target: Union[ast.expr, list]) -> bool: # type: ignore[type-arg] - """Check if the assignment target is the variable 'codeflash_output'.""" - if isinstance(target, ast.Name): - return target.id == "codeflash_output" - if isinstance(target, (ast.Tuple, ast.List)): - # Handle tuple/list unpacking: a, codeflash_output, b = values - return any(self._is_codeflash_output_target(elt) for elt in target.elts) - if isinstance(target, (ast.Subscript, ast.Attribute)): - # Not a simple variable assignment - return False - return False - - def _record_assignment(self, node: ast.AST) -> None: - """Record an assignment to codeflash_output.""" - relative_line = node.lineno - 1 # type: ignore[attr-defined] - self.results.append(relative_line) - - def visit_Assign(self, node: ast.Assign) -> None: - """Visit assignment statements: codeflash_output = value.""" - for target in node.targets: - if self._is_codeflash_output_target(target): - self._record_assignment(node) - break + def visit_Call(self, node): # type: ignore[no-untyped-def] # noqa: ANN201, ANN001 + """Detect fn calls.""" + func_name = self._get_called_func_name(node.func) # type: ignore[no-untyped-call] + if func_name == self.name: + self.results.append(node.lineno - 1) self.generic_visit(node) - def visit_AnnAssign(self, node: ast.AnnAssign) -> None: - """Visit annotated assignments: codeflash_output: int = value.""" - if self._is_codeflash_output_target(node.target): - self._record_assignment(node) - self.generic_visit(node) - - def visit_AugAssign(self, node: ast.AugAssign) -> None: - """Visit augmented assignments: codeflash_output += value.""" - if self._is_codeflash_output_target(node.target): - self._record_assignment(node) - self.generic_visit(node) - - def visit_NamedExpr(self, node: ast.NamedExpr) -> None: - """Visit walrus operator: (codeflash_output := value).""" - if isinstance(node.target, ast.Name) and node.target.id == "codeflash_output": - self._record_assignment(node) - self.generic_visit(node) - - def visit_For(self, node: ast.For) -> None: - """Visit for loops: for codeflash_output in iterable.""" - if self._is_codeflash_output_target(node.target): - self._record_assignment(node) - self.generic_visit(node) - - def visit_comprehension(self, node: ast.comprehension) -> None: - """Visit comprehensions: [x for codeflash_output in iterable].""" - if self._is_codeflash_output_target(node.target): - # Comprehensions don't have line numbers, so we skip recording - pass - self.generic_visit(node) - - def visit_With(self, node: ast.With) -> None: - """Visit with statements: with expr as codeflash_output.""" - for item in node.items: - if item.optional_vars and self._is_codeflash_output_target(item.optional_vars): - self._record_assignment(node) - break - self.generic_visit(node) - - def visit_ExceptHandler(self, node: ast.ExceptHandler) -> None: - """Visit except handlers: except Exception as codeflash_output.""" - if node.name == "codeflash_output": - self._record_assignment(node) - self.generic_visit(node) + def _get_called_func_name(self, node): # type: ignore[no-untyped-def] # noqa: ANN001, ANN202 + """Return name of called fn.""" + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None -def find_codeflash_output_assignments(source_code: str) -> list[int]: +def find_codeflash_output_assignments(function_name: str, source_code: str) -> list[int]: tree = ast.parse(source_code) - visitor = CfoVisitor(source_code) + visitor = CfoVisitor(function_name, source_code) visitor.visit(tree) return visitor.results -def add_runtime_comments_to_generated_tests( - test_cfg: TestConfig, - generated_tests: GeneratedTestsList, - original_runtimes: dict[InvocationId, list[int]], - optimized_runtimes: dict[InvocationId, list[int]], -) -> GeneratedTestsList: - """Add runtime performance comments to function calls in generated tests.""" - tests_root = test_cfg.tests_root - module_root = test_cfg.project_root_path - rel_tests_root = tests_root.relative_to(module_root) - - # TODO: reduce for loops to one - class RuntimeCommentTransformer(cst.CSTTransformer): - def __init__(self, module: cst.Module, test: GeneratedTests, tests_root: Path, rel_tests_root: Path) -> None: - super().__init__() - self.test = test - self.context_stack: list[str] = [] - self.tests_root = tests_root - self.rel_tests_root = rel_tests_root - self.module = module - self.cfo_locs: list[int] = [] - self.cfo_idx_loc_to_look_at: int = -1 - - def visit_ClassDef(self, node: cst.ClassDef) -> None: - # Track when we enter a class - self.context_stack.append(node.name.value) - - def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: # noqa: ARG002 - # Pop the context when we leave a class - self.context_stack.pop() - return updated_node - - def visit_FunctionDef(self, node: cst.FunctionDef) -> None: - # convert function body to ast normalized string and find occurrences of codeflash_output - body_code = dedent(self.module.code_for_node(node.body)) - normalized_body_code = ast.unparse(ast.parse(body_code)) - self.cfo_locs = sorted( - find_codeflash_output_assignments(normalized_body_code) - ) # sorted in order we will encounter them - self.cfo_idx_loc_to_look_at = -1 - self.context_stack.append(node.name.value) - - def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: # noqa: ARG002 - # Pop the context when we leave a function - self.context_stack.pop() - return updated_node - - def leave_SimpleStatementLine( - self, - original_node: cst.SimpleStatementLine, # noqa: ARG002 - updated_node: cst.SimpleStatementLine, - ) -> cst.SimpleStatementLine: - # Look for assignment statements that assign to codeflash_output - # Handle both single statements and multiple statements on one line - codeflash_assignment_found = False - for stmt in updated_node.body: - if isinstance(stmt, cst.Assign) and ( - len(stmt.targets) == 1 - and isinstance(stmt.targets[0].target, cst.Name) - and stmt.targets[0].target.value == "codeflash_output" +class Finder(cst.CSTVisitor): + def __init__(self, name: str) -> None: + super().__init__() + self.found = False + self.name = name + + def visit_Call(self, call_node) -> None: # type: ignore[no-untyped-def] # noqa : ANN001 + func_expr = call_node.func + if isinstance(func_expr, cst.Name): + if func_expr.value == self.name: + self.found = True + elif isinstance(func_expr, cst.Attribute): # noqa : SIM102 + if func_expr.attr.value == self.name: + self.found = True + + +# TODO: reduce for loops to one +class RuntimeCommentTransformer(cst.CSTTransformer): + def __init__( + self, + qualified_name: str, + module: cst.Module, + test: GeneratedTests, + tests_root: Path, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], + ) -> None: + super().__init__() + self.test = test + self.context_stack: list[str] = [] + self.tests_root = tests_root + self.module = module + self.cfo_locs: list[int] = [] + self.cfo_idx_loc_to_look_at: int = -1 + self.name = qualified_name.split(".")[-1] + self.original_runtimes = original_runtimes + self.optimized_runtimes = optimized_runtimes + + def visit_ClassDef(self, node: cst.ClassDef) -> None: + # Track when we enter a class + self.context_stack.append(node.name.value) + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: # noqa: ARG002 + # Pop the context when we leave a class + self.context_stack.pop() + return updated_node + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + # convert function body to ast normalized string and find occurrences of codeflash_output + body_code = dedent(self.module.code_for_node(node.body)) + normalized_body_code = ast.unparse(ast.parse(body_code)) + self.cfo_locs = sorted( + find_codeflash_output_assignments(self.name, normalized_body_code) + ) # sorted in order we will encounter them + self.cfo_idx_loc_to_look_at = -1 + self.context_stack.append(node.name.value) + + def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: # noqa: ARG002 + # Pop the context when we leave a function + self.context_stack.pop() + return updated_node + + def leave_SimpleStatementLine( + self, + original_node: cst.SimpleStatementLine, # noqa: ARG002 + updated_node: cst.SimpleStatementLine, + ) -> cst.SimpleStatementLine: + # Check if this statement line contains a call to self.name + if self._contains_myfunc_call(updated_node): # type: ignore[no-untyped-call] + # Find matching test cases by looking for this test function name in the test results + self.cfo_idx_loc_to_look_at += 1 + matching_original_times = [] + matching_optimized_times = [] + # TODO : will not work if there are multiple test cases with the same name, match filename + test class + test function name + invocationid + for invocation_id, runtimes in self.original_runtimes.items(): + # get position here and match in if condition + qualified_name = ( + invocation_id.test_class_name + "." + invocation_id.test_function_name # type: ignore[operator] + if invocation_id.test_class_name + else invocation_id.test_function_name + ) + abs_path = Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve() + if ( + qualified_name == ".".join(self.context_stack) + and abs_path in [self.test.behavior_file_path, self.test.perf_file_path] + and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at] # type:ignore[union-attr] ): - codeflash_assignment_found = True - break - - if codeflash_assignment_found: - # Find matching test cases by looking for this test function name in the test results - self.cfo_idx_loc_to_look_at += 1 - matching_original_times = [] - matching_optimized_times = [] - # TODO : will not work if there are multiple test cases with the same name, match filename + test class + test function name + invocationid - for invocation_id, runtimes in original_runtimes.items(): - # get position here and match in if condition - qualified_name = ( - invocation_id.test_class_name + "." + invocation_id.test_function_name # type: ignore[operator] - if invocation_id.test_class_name - else invocation_id.test_function_name - ) - rel_path = ( - Path(invocation_id.test_module_path.replace(".", os.sep)) - .with_suffix(".py") - .relative_to(self.rel_tests_root) - ) - if ( - qualified_name == ".".join(self.context_stack) - and rel_path - in [ - self.test.behavior_file_path.relative_to(self.tests_root), - self.test.perf_file_path.relative_to(self.tests_root), - ] - and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at] # type:ignore[union-attr] - ): - matching_original_times.extend(runtimes) - - for invocation_id, runtimes in optimized_runtimes.items(): - # get position here and match in if condition - qualified_name = ( - invocation_id.test_class_name + "." + invocation_id.test_function_name # type: ignore[operator] - if invocation_id.test_class_name - else invocation_id.test_function_name + matching_original_times.extend(runtimes) + + for invocation_id, runtimes in self.optimized_runtimes.items(): + # get position here and match in if condition + qualified_name = ( + invocation_id.test_class_name + "." + invocation_id.test_function_name # type: ignore[operator] + if invocation_id.test_class_name + else invocation_id.test_function_name + ) + abs_path = Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve() + if ( + qualified_name == ".".join(self.context_stack) + and abs_path in [self.test.behavior_file_path, self.test.perf_file_path] + and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at] # type:ignore[union-attr] + ): + matching_optimized_times.extend(runtimes) + + if matching_original_times and matching_optimized_times: + original_time = min(matching_original_times) + optimized_time = min(matching_optimized_times) + if original_time != 0 and optimized_time != 0: + perf_gain = format_perf( + abs( + performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) + * 100 + ) ) - rel_path = ( - Path(invocation_id.test_module_path.replace(".", os.sep)) - .with_suffix(".py") - .relative_to(self.rel_tests_root) + status = "slower" if optimized_time > original_time else "faster" + # Create the runtime comment + comment_text = ( + f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})" ) - if ( - qualified_name == ".".join(self.context_stack) - and rel_path - in [ - self.test.behavior_file_path.relative_to(self.tests_root), - self.test.perf_file_path.relative_to(self.tests_root), - ] - and int(invocation_id.iteration_id.split("_")[0]) == self.cfo_locs[self.cfo_idx_loc_to_look_at] # type:ignore[union-attr] - ): - matching_optimized_times.extend(runtimes) - - if matching_original_times and matching_optimized_times: - original_time = min(matching_original_times) - optimized_time = min(matching_optimized_times) - if original_time != 0 and optimized_time != 0: - perf_gain = format_perf( - abs( - performance_gain(original_runtime_ns=original_time, optimized_runtime_ns=optimized_time) - * 100 - ) - ) - status = "slower" if optimized_time > original_time else "faster" - # Create the runtime comment - comment_text = ( - f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})" - ) - - # Add comment to the trailing whitespace - new_trailing_whitespace = cst.TrailingWhitespace( + return updated_node.with_changes( + trailing_whitespace=cst.TrailingWhitespace( whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment(comment_text), newline=updated_node.trailing_whitespace.newline, ) + ) + return updated_node - return updated_node.with_changes(trailing_whitespace=new_trailing_whitespace) + def _contains_myfunc_call(self, node): # type: ignore[no-untyped-def] # noqa : ANN202, ANN001 + """Recursively search for any Call node in the statement whose function is named self.name (including obj.myfunc).""" + finder = Finder(self.name) + node.visit(finder) + return finder.found - return updated_node + +def add_runtime_comments_to_generated_tests( + qualified_name: str, + test_cfg: TestConfig, + generated_tests: GeneratedTestsList, + original_runtimes: dict[InvocationId, list[int]], + optimized_runtimes: dict[InvocationId, list[int]], +) -> GeneratedTestsList: + """Add runtime performance comments to function calls in generated tests.""" + tests_root = test_cfg.tests_root # Process each generated test modified_tests = [] @@ -282,7 +230,10 @@ def leave_SimpleStatementLine( # Parse the test source code tree = cst.parse_module(test.generated_original_test_source) # Transform the tree to add runtime comments - transformer = RuntimeCommentTransformer(tree, test, tests_root, rel_tests_root) + # qualified_name: str, module: cst.Module, test: GeneratedTests, tests_root: Path + transformer = RuntimeCommentTransformer( + qualified_name, tree, test, tests_root, original_runtimes, optimized_runtimes + ) modified_tree = tree.visit(transformer) # Convert back to source code diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 82cf4bc57..8772bb9e5 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1012,15 +1012,20 @@ def find_and_process_best_optimization( optimized_runtime_by_test = ( best_optimization.winning_benchmarking_test_results.usable_runtime_data_by_test_case() ) + qualified_name = self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root) # Add runtime comments to generated tests before creating the PR generated_tests = add_runtime_comments_to_generated_tests( - self.test_cfg, generated_tests, original_runtime_by_test, optimized_runtime_by_test + qualified_name, + self.test_cfg, + generated_tests, + original_runtime_by_test, + optimized_runtime_by_test, ) generated_tests_str = "\n\n".join( [test.generated_original_test_source for test in generated_tests.generated_tests] ) existing_tests = existing_tests_source_for( - self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root), + qualified_name, function_to_all_tests, test_cfg=self.test_cfg, original_runtimes_all=original_runtime_by_test, diff --git a/codeflash/result/create_pr.py b/codeflash/result/create_pr.py index a08875a4f..ee8e37635 100644 --- a/codeflash/result/create_pr.py +++ b/codeflash/result/create_pr.py @@ -42,43 +42,39 @@ def existing_tests_source_for( rows = [] headers = ["Test File::Test Function", "Original ⏱️", "Optimized ⏱️", "Speedup"] tests_root = test_cfg.tests_root - module_root = test_cfg.project_root_path - rel_tests_root = tests_root.relative_to(module_root) original_tests_to_runtimes: dict[Path, dict[str, int]] = {} optimized_tests_to_runtimes: dict[Path, dict[str, int]] = {} non_generated_tests = set() for test_file in test_files: - non_generated_tests.add(Path(test_file.tests_in_file.test_file).relative_to(tests_root)) + non_generated_tests.add(test_file.tests_in_file.test_file) # TODO confirm that original and optimized have the same keys all_invocation_ids = original_runtimes_all.keys() | optimized_runtimes_all.keys() for invocation_id in all_invocation_ids: - rel_path = ( - Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").relative_to(rel_tests_root) - ) - if rel_path not in non_generated_tests: + abs_path = Path(invocation_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve() + if abs_path not in non_generated_tests: continue - if rel_path not in original_tests_to_runtimes: - original_tests_to_runtimes[rel_path] = {} - if rel_path not in optimized_tests_to_runtimes: - optimized_tests_to_runtimes[rel_path] = {} + if abs_path not in original_tests_to_runtimes: + original_tests_to_runtimes[abs_path] = {} + if abs_path not in optimized_tests_to_runtimes: + optimized_tests_to_runtimes[abs_path] = {} qualified_name = ( invocation_id.test_class_name + "." + invocation_id.test_function_name # type: ignore[operator] if invocation_id.test_class_name else invocation_id.test_function_name ) - if qualified_name not in original_tests_to_runtimes[rel_path]: - original_tests_to_runtimes[rel_path][qualified_name] = 0 # type: ignore[index] - if qualified_name not in optimized_tests_to_runtimes[rel_path]: - optimized_tests_to_runtimes[rel_path][qualified_name] = 0 # type: ignore[index] + if qualified_name not in original_tests_to_runtimes[abs_path]: + original_tests_to_runtimes[abs_path][qualified_name] = 0 # type: ignore[index] + if qualified_name not in optimized_tests_to_runtimes[abs_path]: + optimized_tests_to_runtimes[abs_path][qualified_name] = 0 # type: ignore[index] if invocation_id in original_runtimes_all: - original_tests_to_runtimes[rel_path][qualified_name] += min(original_runtimes_all[invocation_id]) # type: ignore[index] + original_tests_to_runtimes[abs_path][qualified_name] += min(original_runtimes_all[invocation_id]) # type: ignore[index] if invocation_id in optimized_runtimes_all: - optimized_tests_to_runtimes[rel_path][qualified_name] += min(optimized_runtimes_all[invocation_id]) # type: ignore[index] + optimized_tests_to_runtimes[abs_path][qualified_name] += min(optimized_runtimes_all[invocation_id]) # type: ignore[index] # parse into string - all_rel_paths = ( + all_abs_paths = ( original_tests_to_runtimes.keys() ) # both will have the same keys as some default values are assigned in the previous loop - for filename in sorted(all_rel_paths): + for filename in sorted(all_abs_paths): all_qualified_names = original_tests_to_runtimes[ filename ].keys() # both will have the same keys as some default values are assigned in the previous loop @@ -90,6 +86,7 @@ def existing_tests_source_for( ): print_optimized_runtime = format_time(optimized_tests_to_runtimes[filename][qualified_name]) print_original_runtime = format_time(original_tests_to_runtimes[filename][qualified_name]) + print_filename = filename.relative_to(tests_root) greater = ( optimized_tests_to_runtimes[filename][qualified_name] > original_tests_to_runtimes[filename][qualified_name] @@ -104,7 +101,7 @@ def existing_tests_source_for( if greater: rows.append( [ - f"`{filename}::{qualified_name}`", + f"`{print_filename}::{qualified_name}`", f"{print_original_runtime}", f"{print_optimized_runtime}", f"⚠️{perf_gain}%", @@ -113,7 +110,7 @@ def existing_tests_source_for( else: rows.append( [ - f"`{filename}::{qualified_name}`", + f"`{print_filename}::{qualified_name}`", f"{print_original_runtime}", f"{print_optimized_runtime}", f"✅{perf_gain}%", diff --git a/tests/test_add_runtime_comments.py b/tests/test_add_runtime_comments.py index 71f1d7566..d2f31e724 100644 --- a/tests/test_add_runtime_comments.py +++ b/tests/test_add_runtime_comments.py @@ -9,14 +9,19 @@ VerificationType, TestResults from codeflash.verification.verification_utils import TestConfig + +TestType.__test__ = False +TestConfig.__test__ = False +TestResults.__test__ = False + @pytest.fixture def test_config(): """Create a mock TestConfig for testing.""" config = Mock(spec=TestConfig) - config.project_root_path = Path("/project") + config.project_root_path = Path(__file__).resolve().parent.parent config.test_framework= "pytest" - config.tests_project_rootdir = Path("/project/tests") - config.tests_root = Path("/project/tests") + config.tests_project_rootdir = Path(__file__).resolve().parent + config.tests_root = Path(__file__).resolve().parent return config class TestAddRuntimeComments: @@ -48,17 +53,19 @@ def create_test_invocation( def test_basic_runtime_comment_addition(self, test_config): """Test basic functionality of adding runtime comments.""" # Create test source code + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ + qualified_name = "bubble_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py", ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -75,7 +82,7 @@ def test_basic_runtime_comment_addition(self, test_config): original_runtimes = original_test_results.usable_runtime_data_by_test_case() optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that comments were added modified_source = result.generated_tests[0].generated_original_test_source @@ -84,8 +91,9 @@ def test_basic_runtime_comment_addition(self, test_config): def test_multiple_test_functions(self, test_config): """Test handling multiple test functions in the same file.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): - codeflash_output = bubble_sort([3, 1, 2]) + codeflash_output = quick_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] def test_quick_sort(): @@ -95,13 +103,13 @@ def test_quick_sort(): def helper_function(): return "not a test" """ - + qualified_name = "quick_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -121,7 +129,7 @@ def helper_function(): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) modified_source = result.generated_tests[0].generated_original_test_source @@ -136,6 +144,7 @@ def helper_function(): def test_different_time_formats(self, test_config): """Test that different time ranges are formatted correctly with new precision rules.""" + os.chdir(test_config.project_root_path) test_cases = [ (999, 500, "999ns -> 500ns"), # nanoseconds (25_000, 18_000, "25.0μs -> 18.0μs"), # microseconds with precision @@ -151,13 +160,13 @@ def test_different_time_formats(self, test_config): codeflash_output = some_function() assert codeflash_output is not None """ - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -173,7 +182,7 @@ def test_different_time_formats(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) modified_source = result.generated_tests[0].generated_original_test_source @@ -181,17 +190,19 @@ def test_different_time_formats(self, test_config): def test_missing_test_results(self, test_config): """Test behavior when test results are missing for a test function.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ + qualified_name = "bubble_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -204,7 +215,7 @@ def test_missing_test_results(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that no comments were added modified_source = result.generated_tests[0].generated_original_test_source @@ -212,17 +223,19 @@ def test_missing_test_results(self, test_config): def test_partial_test_results(self, test_config): """Test behavior when only one set of test results is available.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ + qualified_name = "bubble_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -236,7 +249,7 @@ def test_partial_test_results(self, test_config): original_runtimes = original_test_results.usable_runtime_data_by_test_case() optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that no comments were added modified_source = result.generated_tests[0].generated_original_test_source @@ -244,17 +257,18 @@ def test_partial_test_results(self, test_config): def test_multiple_runtimes_uses_minimum(self, test_config): """Test that when multiple runtimes exist, the minimum is used.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ - + qualified_name = "bubble_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -275,7 +289,7 @@ def test_multiple_runtimes_uses_minimum(self, test_config): original_runtimes = original_test_results.usable_runtime_data_by_test_case() optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that minimum times were used (500μs -> 300μs) modified_source = result.generated_tests[0].generated_original_test_source @@ -283,17 +297,18 @@ def test_multiple_runtimes_uses_minimum(self, test_config): def test_no_codeflash_output_assignment(self, test_config): """Test behavior when test doesn't have codeflash_output assignment.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): result = bubble_sort([3, 1, 2]) assert result == [1, 2, 3] """ - + qualified_name = "bubble_sort" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -309,7 +324,7 @@ def test_no_codeflash_output_assignment(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that no comments were added (no codeflash_output assignment) modified_source = result.generated_tests[0].generated_original_test_source @@ -317,6 +332,7 @@ def test_no_codeflash_output_assignment(self, test_config): def test_invalid_python_code_handling(self, test_config): """Test behavior when test source code is invalid Python.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(: codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] @@ -326,10 +342,10 @@ def test_invalid_python_code_handling(self, test_config): generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) - + qualified_name = "bubble_sort" generated_tests = GeneratedTestsList(generated_tests=[generated_test]) # Create test results @@ -343,7 +359,7 @@ def test_invalid_python_code_handling(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - should handle parse error gracefully - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that original test is preserved when parsing fails modified_source = result.generated_tests[0].generated_original_test_source @@ -351,8 +367,9 @@ def test_invalid_python_code_handling(self, test_config): def test_multiple_generated_tests(self, test_config): """Test handling multiple generated test objects.""" + os.chdir(test_config.project_root_path) test_source_1 = """def test_bubble_sort(): - codeflash_output = bubble_sort([3, 1, 2]) + codeflash_output = quick_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ @@ -363,21 +380,21 @@ def test_multiple_generated_tests(self, test_config): codeflash_output = quick_sort([5, 2, 8]) assert codeflash_output == [2, 5, 8] """ - + qualified_name = "quick_sort" generated_test_1 = GeneratedTests( generated_original_test_source=test_source_1, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_test_2 = GeneratedTests( generated_original_test_source=test_source_2, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test_1, generated_test_2]) @@ -396,7 +413,7 @@ def test_multiple_generated_tests(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that comments were added to both test files modified_source_1 = result.generated_tests[0].generated_original_test_source @@ -407,22 +424,23 @@ def test_multiple_generated_tests(self, test_config): def test_preserved_test_attributes(self, test_config): """Test that other test attributes are preserved during modification.""" + os.chdir(test_config.project_root_path) test_source = """def test_bubble_sort(): codeflash_output = bubble_sort([3, 1, 2]) assert codeflash_output == [1, 2, 3] """ - + qualified_name = "bubble_sort" original_behavior_source = "behavior test source" original_perf_source = "perf test source" - original_behavior_path = Path("/project/tests/test_module.py") - original_perf_path = Path("/project/tests/test_module_perf.py") + original_behavior_path=test_config.tests_root / "test_module.py" + original_perf_path=test_config.tests_root / "test_perf.py" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source=original_behavior_source, instrumented_perf_test_source=original_perf_source, - behavior_file_path=original_behavior_path, - perf_file_path=original_perf_path + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -437,7 +455,7 @@ def test_preserved_test_attributes(self, test_config): original_runtimes = original_test_results.usable_runtime_data_by_test_case() optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that other attributes are preserved modified_test = result.generated_tests[0] @@ -451,6 +469,7 @@ def test_preserved_test_attributes(self, test_config): def test_multistatement_line_handling(self, test_config): """Test that runtime comments work correctly with multiple statements on one line.""" + os.chdir(test_config.project_root_path) test_source = """def test_mutation_of_input(): # Test that the input list is mutated in-place and returned arr = [3, 1, 2] @@ -458,13 +477,13 @@ def test_multistatement_line_handling(self, test_config): assert result == [1, 2, 3] assert arr == [1, 2, 3] # Input should be mutated """ - + qualified_name = "sorter" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py") + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -480,7 +499,7 @@ def test_multistatement_line_handling(self, test_config): optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() # Test the functionality - result = add_runtime_comments_to_generated_tests(test_config, generated_tests, original_runtimes, optimized_runtimes) + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) # Check that comments were added to the correct line modified_source = result.generated_tests[0].generated_original_test_source @@ -500,17 +519,18 @@ def test_multistatement_line_handling(self, test_config): def test_add_runtime_comments_simple_function(self, test_config): """Test adding runtime comments to a simple test function.""" + os.chdir(test_config.project_root_path) test_source = '''def test_function(): codeflash_output = some_function() assert codeflash_output == expected ''' - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -527,7 +547,7 @@ def test_add_runtime_comments_simple_function(self, test_config): optimized_runtimes = {invocation_id: [500000000, 600000000]} # 0.5s, 0.6s in nanoseconds result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) expected_source = '''def test_function(): @@ -540,18 +560,19 @@ def test_add_runtime_comments_simple_function(self, test_config): def test_add_runtime_comments_class_method(self, test_config): """Test adding runtime comments to a test method within a class.""" + os.chdir(test_config.project_root_path) test_source = '''class TestClass: def test_function(self): codeflash_output = some_function() assert codeflash_output == expected ''' - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -569,7 +590,7 @@ def test_function(self): optimized_runtimes = {invocation_id: [1000000000]} # 1s in nanoseconds result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) expected_source = '''class TestClass: @@ -583,20 +604,23 @@ def test_function(self): def test_add_runtime_comments_multiple_assignments(self, test_config): """Test adding runtime comments when there are multiple codeflash_output assignments.""" + os.chdir(test_config.project_root_path) test_source = '''def test_function(): setup_data = prepare_test() codeflash_output = some_function() assert codeflash_output == expected codeflash_output = another_function() assert codeflash_output == expected2 + codeflash_output = some_function() + assert codeflash_output == expected2 ''' - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -612,22 +636,24 @@ def test_add_runtime_comments_multiple_assignments(self, test_config): test_module_path="tests.test_module", test_class_name=None, test_function_name="test_function", - function_getting_tested="another_function", - iteration_id="3", + function_getting_tested="some_function", + iteration_id="5", ) original_runtimes = {invocation_id1: [1500000000], invocation_id2: [10]} # 1.5s in nanoseconds optimized_runtimes = {invocation_id1: [750000000], invocation_id2: [5]} # 0.75s in nanoseconds result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) expected_source = '''def test_function(): setup_data = prepare_test() codeflash_output = some_function() # 1.50s -> 750ms (100% faster) assert codeflash_output == expected - codeflash_output = another_function() # 10ns -> 5ns (100% faster) + codeflash_output = another_function() + assert codeflash_output == expected2 + codeflash_output = some_function() # 10ns -> 5ns (100% faster) assert codeflash_output == expected2 ''' @@ -636,17 +662,18 @@ def test_add_runtime_comments_multiple_assignments(self, test_config): def test_add_runtime_comments_no_matching_runtimes(self, test_config): """Test that source remains unchanged when no matching runtimes are found.""" + os.chdir(test_config.project_root_path) test_source = '''def test_function(): codeflash_output = some_function() assert codeflash_output == expected ''' - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -664,7 +691,7 @@ def test_add_runtime_comments_no_matching_runtimes(self, test_config): optimized_runtimes = {invocation_id: [500000000]} result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) # Source should remain unchanged @@ -672,18 +699,23 @@ def test_add_runtime_comments_no_matching_runtimes(self, test_config): assert result.generated_tests[0].generated_original_test_source == test_source def test_add_runtime_comments_no_codeflash_output(self, test_config): - """Test that source remains unchanged when there's no codeflash_output assignment.""" + """comments will still be added if codeflash output doesnt exist""" + os.chdir(test_config.project_root_path) test_source = '''def test_function(): result = some_function() assert result == expected ''' - + qualified_name = "some_function" + expected = '''def test_function(): + result = some_function() # 1.00s -> 500ms (100% faster) + assert result == expected +''' generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -700,39 +732,40 @@ def test_add_runtime_comments_no_codeflash_output(self, test_config): optimized_runtimes = {invocation_id: [500000000]} result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) # Source should remain unchanged assert len(result.generated_tests) == 1 - assert result.generated_tests[0].generated_original_test_source == test_source + assert result.generated_tests[0].generated_original_test_source == expected def test_add_runtime_comments_multiple_tests(self, test_config): """Test adding runtime comments to multiple generated tests.""" + os.chdir(test_config.project_root_path) test_source1 = '''def test_function1(): codeflash_output = some_function() assert codeflash_output == expected ''' test_source2 = '''def test_function2(): - codeflash_output = another_function() + codeflash_output = some_function() assert codeflash_output == expected ''' - + qualified_name = "some_function" generated_test1 = GeneratedTests( generated_original_test_source=test_source1, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module1.py"), - perf_file_path=Path("/project/tests/test_module1_perf.py"), + behavior_file_path=test_config.tests_root / "test_module1.py", + perf_file_path=test_config.tests_root / "test_perf1.py" ) generated_test2 = GeneratedTests( generated_original_test_source=test_source2, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module2.py"), - perf_file_path=Path("/project/tests/test_module2_perf.py"), + behavior_file_path=test_config.tests_root / "test_module2.py", + perf_file_path=test_config.tests_root / "test_perf2.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test1, generated_test2]) @@ -749,7 +782,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config): test_module_path="tests.test_module2", test_class_name=None, test_function_name="test_function2", - function_getting_tested="another_function", + function_getting_tested="some_function", # not used in this test throughout the entire test file iteration_id = "0", ) @@ -763,7 +796,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config): } result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) expected_source1 = '''def test_function1(): @@ -772,7 +805,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config): ''' expected_source2 = '''def test_function2(): - codeflash_output = another_function() # 2.00s -> 800ms (150% faster) + codeflash_output = some_function() # 2.00s -> 800ms (150% faster) assert codeflash_output == expected ''' @@ -782,19 +815,20 @@ def test_add_runtime_comments_multiple_tests(self, test_config): def test_add_runtime_comments_performance_regression(self, test_config): """Test adding runtime comments when optimized version is slower (negative performance gain).""" + os.chdir(test_config.project_root_path) test_source = '''def test_function(): codeflash_output = some_function() assert codeflash_output == expected codeflash_output = some_function() assert codeflash_output == expected ''' - + qualified_name = "some_function" generated_test = GeneratedTests( generated_original_test_source=test_source, instrumented_behavior_test_source="", instrumented_perf_test_source="", - behavior_file_path=Path("/project/tests/test_module.py"), - perf_file_path=Path("/project/tests/test_module_perf.py"), + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" ) generated_tests = GeneratedTestsList(generated_tests=[generated_test]) @@ -819,7 +853,7 @@ def test_add_runtime_comments_performance_regression(self, test_config): optimized_runtimes = {invocation_id1: [1500000000], invocation_id2: [1]} # 1.5s (slower!) result = add_runtime_comments_to_generated_tests( - test_config, generated_tests, original_runtimes, optimized_runtimes + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes ) expected_source = '''def test_function(): @@ -831,3 +865,769 @@ def test_add_runtime_comments_performance_regression(self, test_config): assert len(result.generated_tests) == 1 assert result.generated_tests[0].generated_original_test_source == expected_source + + def test_basic_runtime_comment_addition_no_cfo(self, test_config): + """Test basic functionality of adding runtime comments.""" + # Create test source code + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + + qualified_name = "bubble_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py", + ) + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + # Add test invocations with different runtimes + original_invocation = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='0') # 500μs + optimized_invocation = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='0') # 300μs + + original_test_results.add(original_invocation) + optimized_test_results.add(optimized_invocation) + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that comments were added + modified_source = result.generated_tests[0].generated_original_test_source + assert "# 500μs -> 300μs" in modified_source + assert "result = bubble_sort([3, 1, 2]) # 500μs -> 300μs" in modified_source + + def test_multiple_test_functions_no_cfo(self, test_config): + """Test handling multiple test functions in the same file.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = quick_sort([3, 1, 2]) + assert result == [1, 2, 3] + +def test_quick_sort(): + result = quick_sort([5, 2, 8]); assert result == [2, 5, 8] + +def helper_function(): + return "not a test" +""" + qualified_name = "quick_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results for both functions + original_test_results = TestResults() + optimized_test_results = TestResults() + + # Add test invocations for both test functions + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='0')) + original_test_results.add(self.create_test_invocation("test_quick_sort", 800_000, iteration_id='0')) + + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_quick_sort", 600_000, iteration_id='0')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + modified_source = result.generated_tests[0].generated_original_test_source + + # Check that comments were added to both test functions + assert "# 500μs -> 300μs" in modified_source + assert "# 800μs -> 600μs" in modified_source + # Helper function should not have comments + assert ( + "helper_function():" in modified_source + and "# " not in modified_source.split("helper_function():")[1].split("\n")[0] + ) + + def test_different_time_formats_no_cfo(self, test_config): + """Test that different time ranges are formatted correctly with new precision rules.""" + os.chdir(test_config.project_root_path) + test_cases = [ + (999, 500, "999ns -> 500ns"), # nanoseconds + (25_000, 18_000, "25.0μs -> 18.0μs"), # microseconds with precision + (500_000, 300_000, "500μs -> 300μs"), # microseconds full integers + (1_500_000, 800_000, "1.50ms -> 800μs"), # milliseconds with precision + (365_000_000, 290_000_000, "365ms -> 290ms"), # milliseconds full integers + (2_000_000_000, 1_500_000_000, "2.00s -> 1.50s"), # seconds with precision + ] + + for original_time, optimized_time, expected_comment in test_cases: + test_source = """def test_function(): + #this comment will be removed in ast form + result = some_function(); assert result is not None +""" + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_function", original_time, iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_function", optimized_time, iteration_id='0')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + # Test the functionality + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + modified_source = result.generated_tests[0].generated_original_test_source + assert f"# {expected_comment}" in modified_source + + def test_missing_test_results_no_cfo(self, test_config): + """Test behavior when test results are missing for a test function.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + + qualified_name = "bubble_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create empty test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that no comments were added + modified_source = result.generated_tests[0].generated_original_test_source + assert modified_source == test_source # Should be unchanged + + def test_partial_test_results_no_cfo(self, test_config): + """Test behavior when only one set of test results is available.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + + qualified_name = "bubble_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results with only original data + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='0')) + # No optimized results + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that no comments were added + modified_source = result.generated_tests[0].generated_original_test_source + assert modified_source == test_source # Should be unchanged + + def test_multiple_runtimes_uses_minimum_no_cfo(self, test_config): + """Test that when multiple runtimes exist, the minimum is used.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + qualified_name = "bubble_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results with multiple loop iterations + original_test_results = TestResults() + optimized_test_results = TestResults() + + # Add multiple runs with different runtimes + original_test_results.add(self.create_test_invocation("test_bubble_sort", 600_000, loop_index=1,iteration_id='0')) + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000, loop_index=2,iteration_id='0')) + original_test_results.add(self.create_test_invocation("test_bubble_sort", 550_000, loop_index=3,iteration_id='0')) + + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 350_000, loop_index=1,iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000, loop_index=2,iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 320_000, loop_index=3,iteration_id='0')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that minimum times were used (500μs -> 300μs) + modified_source = result.generated_tests[0].generated_original_test_source + assert "# 500μs -> 300μs" in modified_source + + def test_no_codeflash_output_assignment_invalid_iteration_id(self, test_config): + """Test behavior when test doesn't have codeflash_output assignment.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + qualified_name = "bubble_sort" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000,iteration_id='-1')) + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000,iteration_id='-1')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that no comments were added (no codeflash_output assignment) + modified_source = result.generated_tests[0].generated_original_test_source + assert modified_source == test_source # Should be unchanged + + def test_invalid_python_code_handling_no_cfo(self, test_config): + """Test behavior when test source code is invalid Python.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(: + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" # Invalid syntax: extra indentation + + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + qualified_name = "bubble_sort" + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000,iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000,iteration_id='0')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality - should handle parse error gracefully + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that original test is preserved when parsing fails + modified_source = result.generated_tests[0].generated_original_test_source + assert modified_source == test_source # Should be unchanged due to parse error + + def test_multiple_generated_tests_no_cfo(self, test_config): + """Test handling multiple generated test objects.""" + os.chdir(test_config.project_root_path) + test_source_1 = """def test_bubble_sort(): + codeflash_output = quick_sort([3, 1, 2]); assert codeflash_output == [1, 2, 3] +""" + + test_source_2 = """def test_quick_sort(): + a=1 + b=2 + c=3 + result = quick_sort([5, 2, 8]) + assert result == [2, 5, 8] +""" + qualified_name = "quick_sort" + generated_test_1 = GeneratedTests( + generated_original_test_source=test_source_1, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_test_2 = GeneratedTests( + generated_original_test_source=test_source_2, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test_1, generated_test_2]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000,iteration_id='0')) + original_test_results.add(self.create_test_invocation("test_quick_sort", 800_000,iteration_id='3')) + + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000,iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_quick_sort", 600_000,iteration_id='3')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that comments were added to both test files + modified_source_1 = result.generated_tests[0].generated_original_test_source + modified_source_2 = result.generated_tests[1].generated_original_test_source + + assert "# 500μs -> 300μs" in modified_source_1 + assert "# 800μs -> 600μs" in modified_source_2 + + def test_preserved_test_attributes_no_cfo(self, test_config): + """Test that other test attributes are preserved during modification.""" + os.chdir(test_config.project_root_path) + test_source = """def test_bubble_sort(): + result = bubble_sort([3, 1, 2]) + assert result == [1, 2, 3] +""" + qualified_name = "bubble_sort" + original_behavior_source = "behavior test source" + original_perf_source = "perf test source" + original_behavior_path=test_config.tests_root / "test_module.py" + original_perf_path=test_config.tests_root / "test_perf.py" + + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source=original_behavior_source, + instrumented_perf_test_source=original_perf_source, + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_bubble_sort", 500_000,iteration_id='0')) + optimized_test_results.add(self.create_test_invocation("test_bubble_sort", 300_000,iteration_id='0')) + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that other attributes are preserved + modified_test = result.generated_tests[0] + assert modified_test.instrumented_behavior_test_source == original_behavior_source + assert modified_test.instrumented_perf_test_source == original_perf_source + assert modified_test.behavior_file_path == original_behavior_path + assert modified_test.perf_file_path == original_perf_path + + # Check that only the generated_original_test_source was modified + assert "# 500μs -> 300μs" in modified_test.generated_original_test_source + + def test_multistatement_line_handling_no_cfo(self, test_config): + """Test that runtime comments work correctly with multiple statements on one line.""" + os.chdir(test_config.project_root_path) + test_source = """def test_mutation_of_input(): + # Test that the input list is mutated in-place and returned + arr = [3, 1, 2] + res1 = sorter(arr); result = res1 + assert result == [1, 2, 3] + assert arr == [1, 2, 3] # Input should be mutated +""" + qualified_name = "sorter" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Create test results + original_test_results = TestResults() + optimized_test_results = TestResults() + + original_test_results.add(self.create_test_invocation("test_mutation_of_input", 19_000,iteration_id='1')) # 19μs + optimized_test_results.add(self.create_test_invocation("test_mutation_of_input", 14_000,iteration_id='1')) # 14μs + + original_runtimes = original_test_results.usable_runtime_data_by_test_case() + optimized_runtimes = optimized_test_results.usable_runtime_data_by_test_case() + + # Test the functionality + result = add_runtime_comments_to_generated_tests(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes) + + # Check that comments were added to the correct line + modified_source = result.generated_tests[0].generated_original_test_source + assert "# 19.0μs -> 14.0μs" in modified_source + + # Verify the comment is on the line with codeflash_output assignment + lines = modified_source.split("\n") + codeflash_line = None + for line in lines: + if "res1 = sorter(arr)" in line: + codeflash_line = line + break + + assert codeflash_line is not None, "Could not find codeflash_output assignment line" + assert "# 19.0μs -> 14.0μs" in codeflash_line, f"Comment not found in the correct line: {codeflash_line}" + + + def test_add_runtime_comments_simple_function_no_cfo(self, test_config): + """Test adding runtime comments to a simple test function.""" + os.chdir(test_config.project_root_path) + test_source = '''def test_function(): + result = some_function(); assert result == expected +''' + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + invocation_id = InvocationId( + test_module_path="tests.test_module", + test_class_name=None, + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="0", + ) + + original_runtimes = {invocation_id: [1000000000, 1200000000]} # 1s, 1.2s in nanoseconds + optimized_runtimes = {invocation_id: [500000000, 600000000]} # 0.5s, 0.6s in nanoseconds + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + expected_source = '''def test_function(): + result = some_function(); assert result == expected # 1.00s -> 500ms (100% faster) +''' + + assert len(result.generated_tests) == 1 + assert result.generated_tests[0].generated_original_test_source == expected_source + + def test_add_runtime_comments_class_method_no_cfo(self, test_config): + """Test adding runtime comments to a test method within a class.""" + os.chdir(test_config.project_root_path) + test_source = '''class TestClass: + def test_function(self): + result = some_function() + assert result == expected +''' + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + invocation_id = InvocationId( + test_module_path="tests.test_module", + test_class_name="TestClass", + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="0", + + ) + + original_runtimes = {invocation_id: [2000000000]} # 2s in nanoseconds + optimized_runtimes = {invocation_id: [1000000000]} # 1s in nanoseconds + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + expected_source = '''class TestClass: + def test_function(self): + result = some_function() # 2.00s -> 1.00s (100% faster) + assert result == expected +''' + + assert len(result.generated_tests) == 1 + assert result.generated_tests[0].generated_original_test_source == expected_source + + def test_add_runtime_comments_multiple_assignments_no_cfo(self, test_config): + """Test adding runtime comments when there are multiple codeflash_output assignments.""" + os.chdir(test_config.project_root_path) + test_source = '''def test_function(): + setup_data = prepare_test() + codeflash_output = some_function(); assert codeflash_output == expected + result = another_function(); assert result == expected2 + codeflash_output = some_function() + assert codeflash_output == expected2 +''' + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + invocation_id1 = InvocationId( + test_module_path="tests.test_module", + test_class_name=None, + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="1", + ) + invocation_id2 = InvocationId( + test_module_path="tests.test_module", + test_class_name=None, + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="5", + ) + + original_runtimes = {invocation_id1: [1500000000], invocation_id2: [10]} # 1.5s in nanoseconds + optimized_runtimes = {invocation_id1: [750000000], invocation_id2: [5]} # 0.75s in nanoseconds + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + expected_source = '''def test_function(): + setup_data = prepare_test() + codeflash_output = some_function(); assert codeflash_output == expected # 1.50s -> 750ms (100% faster) + result = another_function(); assert result == expected2 + codeflash_output = some_function() # 10ns -> 5ns (100% faster) + assert codeflash_output == expected2 +''' + + assert len(result.generated_tests) == 1 + assert result.generated_tests[0].generated_original_test_source == expected_source + + def test_add_runtime_comments_no_matching_runtimes_no_cfo(self, test_config): + """Test that source remains unchanged when no matching runtimes are found.""" + os.chdir(test_config.project_root_path) + test_source = '''def test_function(): + result = some_function() + assert result == expected +''' + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + # Different invocation ID that won't match + invocation_id = InvocationId( + test_module_path="tests.other_module", + test_class_name=None, + test_function_name="other_function", + function_getting_tested="some_other_function", + iteration_id="0", + ) + + original_runtimes = {invocation_id: [1000000000]} + optimized_runtimes = {invocation_id: [500000000]} + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + # Source should remain unchanged + assert len(result.generated_tests) == 1 + assert result.generated_tests[0].generated_original_test_source == test_source + + + + def test_add_runtime_comments_multiple_tests_no_cfo(self, test_config): + """Test adding runtime comments to multiple generated tests.""" + os.chdir(test_config.project_root_path) + test_source1 = '''def test_function1(): + result = some_function() + assert result == expected +''' + + test_source2 = '''def test_function2(): + result = some_function() + assert result == expected +''' + qualified_name = "some_function" + generated_test1 = GeneratedTests( + generated_original_test_source=test_source1, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module1.py", + perf_file_path=test_config.tests_root / "test_perf1.py" + ) + + generated_test2 = GeneratedTests( + generated_original_test_source=test_source2, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module2.py", + perf_file_path=test_config.tests_root / "test_perf2.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test1, generated_test2]) + + invocation_id1 = InvocationId( + test_module_path="tests.test_module1", + test_class_name=None, + test_function_name="test_function1", + function_getting_tested="some_function", + iteration_id="0", + ) + + invocation_id2 = InvocationId( + test_module_path="tests.test_module2", + test_class_name=None, + test_function_name="test_function2", + function_getting_tested="some_function", # not used in this test throughout the entire test file + iteration_id = "0", + ) + + original_runtimes = { + invocation_id1: [1000000000], # 1s + invocation_id2: [2000000000], # 2s + } + optimized_runtimes = { + invocation_id1: [500000000], # 0.5s + invocation_id2: [800000000], # 0.8s + } + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + expected_source1 = '''def test_function1(): + result = some_function() # 1.00s -> 500ms (100% faster) + assert result == expected +''' + + expected_source2 = '''def test_function2(): + result = some_function() # 2.00s -> 800ms (150% faster) + assert result == expected +''' + + assert len(result.generated_tests) == 2 + assert result.generated_tests[0].generated_original_test_source == expected_source1 + assert result.generated_tests[1].generated_original_test_source == expected_source2 + + def test_add_runtime_comments_performance_regression_no_cfo(self, test_config): + """Test adding runtime comments when optimized version is slower (negative performance gain).""" + os.chdir(test_config.project_root_path) + test_source = '''def test_function(): + result = some_function(); assert codeflash_output == expected + codeflash_output = some_function() + assert codeflash_output == expected +''' + qualified_name = "some_function" + generated_test = GeneratedTests( + generated_original_test_source=test_source, + instrumented_behavior_test_source="", + instrumented_perf_test_source="", + behavior_file_path=test_config.tests_root / "test_module.py", + perf_file_path=test_config.tests_root / "test_perf.py" + ) + + generated_tests = GeneratedTestsList(generated_tests=[generated_test]) + + invocation_id1 = InvocationId( + test_module_path="tests.test_module", + test_class_name=None, + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="0", + ) + + invocation_id2 = InvocationId( + test_module_path="tests.test_module", + test_class_name=None, + test_function_name="test_function", + function_getting_tested="some_function", + iteration_id="2", + ) + + original_runtimes = {invocation_id1: [1000000000], invocation_id2: [2]} # 1s + optimized_runtimes = {invocation_id1: [1500000000], invocation_id2: [1]} # 1.5s (slower!) + + result = add_runtime_comments_to_generated_tests( + qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes + ) + + expected_source = '''def test_function(): + result = some_function(); assert codeflash_output == expected # 1.00s -> 1.50s (33.3% slower) + codeflash_output = some_function() # 2ns -> 1ns (100% faster) + assert codeflash_output == expected +''' + + assert len(result.generated_tests) == 1 + assert result.generated_tests[0].generated_original_test_source == expected_source diff --git a/tests/test_existing_tests_source_for.py b/tests/test_existing_tests_source_for.py index 8940b20d2..ccc3f8ca2 100644 --- a/tests/test_existing_tests_source_for.py +++ b/tests/test_existing_tests_source_for.py @@ -14,8 +14,8 @@ def setup_method(self): """Set up test fixtures.""" # Mock test config self.test_cfg = Mock() - self.test_cfg.tests_root = Path("/project/tests") - self.test_cfg.project_root_path = Path("/project") + self.test_cfg.tests_root = Path(__file__).resolve().parent + self.test_cfg.project_root_path = Path(__file__).resolve().parent.parent # Mock invocation ID self.mock_invocation_id = Mock() @@ -26,10 +26,14 @@ def setup_method(self): # Mock function called in test self.mock_function_called_in_test = Mock() self.mock_function_called_in_test.tests_in_file = Mock() - self.mock_function_called_in_test.tests_in_file.test_file = "/project/tests/test_module.py" + self.mock_function_called_in_test.tests_in_file.test_file = Path(__file__).resolve().parent / "test_module.py" + #Path to pyproject.toml + os.chdir(self.test_cfg.project_root_path) + def test_no_test_files_returns_empty_string(self): """Test that function returns empty string when no test files exist.""" + function_to_tests = {} original_runtimes = {} optimized_runtimes = {} @@ -46,6 +50,7 @@ def test_no_test_files_returns_empty_string(self): def test_single_test_with_improvement(self): """Test single test showing performance improvement.""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -73,6 +78,7 @@ def test_single_test_with_improvement(self): def test_single_test_with_regression(self): """Test single test showing performance regression.""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -100,6 +106,7 @@ def test_single_test_with_regression(self): def test_test_without_class_name(self): """Test function without class name (standalone test function).""" + mock_invocation_no_class = Mock() mock_invocation_no_class.test_module_path = "tests.test_module" mock_invocation_no_class.test_class_name = None @@ -132,6 +139,7 @@ def test_test_without_class_name(self): def test_missing_original_runtime(self): """Test when original runtime is missing (shows NaN).""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -156,6 +164,7 @@ def test_missing_original_runtime(self): def test_missing_optimized_runtime(self): """Test when optimized runtime is missing (shows NaN).""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -181,9 +190,10 @@ def test_missing_optimized_runtime(self): def test_multiple_tests_sorted_output(self): """Test multiple tests with sorted output by filename and function name.""" # Create second test file + mock_function_called_2 = Mock() mock_function_called_2.tests_in_file = Mock() - mock_function_called_2.tests_in_file.test_file = "/project/tests/test_another.py" + mock_function_called_2.tests_in_file.test_file = Path(__file__).resolve().parent / "test_another.py" mock_invocation_2 = Mock() mock_invocation_2.test_module_path = "tests.test_another" @@ -220,6 +230,7 @@ def test_multiple_tests_sorted_output(self): def test_multiple_runtimes_uses_minimum(self): """Test that function uses minimum runtime when multiple measurements exist.""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -247,6 +258,7 @@ def test_multiple_runtimes_uses_minimum(self): def test_complex_module_path_conversion(self): """Test conversion of complex module paths to file paths.""" + mock_invocation_complex = Mock() mock_invocation_complex.test_module_path = "tests.integration.test_complex_module" mock_invocation_complex.test_class_name = "TestComplex" @@ -254,7 +266,7 @@ def test_complex_module_path_conversion(self): mock_function_complex = Mock() mock_function_complex.tests_in_file = Mock() - mock_function_complex.tests_in_file.test_file = f"/project/tests/integration/test_complex_module.py" + mock_function_complex.tests_in_file.test_file = Path(__file__).resolve().parent / "integration/test_complex_module.py" function_to_tests = { "module.function": {mock_function_complex} @@ -283,6 +295,7 @@ def test_complex_module_path_conversion(self): def test_zero_runtime_values(self): """Test handling of zero runtime values.""" + function_to_tests = { "module.function": {self.mock_function_called_in_test} } @@ -310,6 +323,7 @@ def test_zero_runtime_values(self): def test_filters_out_generated_tests(self): """Test that generated tests are filtered out and only non-generated tests are included.""" # Create a test that would be filtered out (not in non_generated_tests) + mock_generated_test = Mock() mock_generated_test.tests_in_file = Mock() mock_generated_test.tests_in_file.test_file = "/project/tests/generated_test.py"