diff --git a/codeflash/code_utils/edit_generated_tests.py b/codeflash/code_utils/edit_generated_tests.py
index a5be1f01a..09c1c163c 100644
--- a/codeflash/code_utils/edit_generated_tests.py
+++ b/codeflash/code_utils/edit_generated_tests.py
@@ -4,10 +4,11 @@
 import os
 import re
 from pathlib import Path
-from textwrap import dedent
 from typing import TYPE_CHECKING
 
 import libcst as cst
+from libcst import MetadataWrapper
+from libcst.metadata import PositionProvider
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.time_utils import format_perf, format_time
@@ -16,230 +17,167 @@
 
 if TYPE_CHECKING:
     from codeflash.models.models import InvocationId
-    from codeflash.verification.verification_utils import TestConfig
 
 
-def remove_functions_from_generated_tests(
-    generated_tests: GeneratedTestsList, test_functions_to_remove: list[str]
-) -> GeneratedTestsList:
-    new_generated_tests = []
-    for generated_test in generated_tests.generated_tests:
-        for test_function in test_functions_to_remove:
-            function_pattern = re.compile(
-                rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?def\s+{re.escape(test_function)}\(.*?\):.*?(?=\ndef\s|$)",
-                re.DOTALL,
-            )
-
-            match = function_pattern.search(generated_test.generated_original_test_source)
-
-            if match is None or "@pytest.mark.parametrize" in match.group(0):
-                continue
-
-            generated_test.generated_original_test_source = function_pattern.sub(
-                "", generated_test.generated_original_test_source
-            )
+class CommentMapper(ast.NodeVisitor):
+    def __init__(
+        self, test: GeneratedTests, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int]
+    ) -> None:
+        self.results: dict[int, str] = {}
+        self.test: GeneratedTests = test
+        self.original_runtimes = original_runtimes
+        self.optimized_runtimes = optimized_runtimes
+        self.abs_path = test.behavior_file_path.with_suffix("")
+        self.context_stack: list[str] = []
 
-        new_generated_tests.append(generated_test)
+    def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
+        self.context_stack.append(node.name)
+        for inner_node in ast.walk(node):
+            if isinstance(inner_node, ast.FunctionDef):
+                self.visit_FunctionDef(inner_node)
+        self.context_stack.pop()
+        return node
+
+    def get_comment(self, match_key: str) -> str:
+        # calculate speedup and output comment
+        original_time = self.original_runtimes[match_key]
+        optimized_time = self.optimized_runtimes[match_key]
+        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
+        return f"# {format_time(original_time)} -> {format_time(optimized_time)} ({perf_gain}% {status})"
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef:
+        self.context_stack.append(node.name)
+        i = len(node.body) - 1
+        test_qualified_name = ".".join(self.context_stack)
+        key = test_qualified_name + "#" + str(self.abs_path)
+        while i >= 0:
+            line_node = node.body[i]
+            if isinstance(line_node, (ast.With, ast.For, ast.While, ast.If)):
+                j = len(line_node.body) - 1
+                while j >= 0:
+                    compound_line_node: ast.stmt = line_node.body[j]
+                    internal_node: ast.AST
+                    for internal_node in ast.walk(compound_line_node):
+                        if isinstance(internal_node, (ast.stmt, ast.Assign)):
+                            inv_id = str(i) + "_" + str(j)
+                            match_key = key + "#" + inv_id
+                            if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
+                                self.results[internal_node.lineno] = self.get_comment(match_key)
+                    j -= 1
+            else:
+                inv_id = str(i)
+                match_key = key + "#" + inv_id
+                if match_key in self.original_runtimes and match_key in self.optimized_runtimes:
+                    self.results[line_node.lineno] = self.get_comment(match_key)
+            i -= 1
+        self.context_stack.pop()
+        return node
 
-    return GeneratedTestsList(generated_tests=new_generated_tests)
 
+def get_fn_call_linenos(
+    test: GeneratedTests, original_runtimes: dict[str, int], optimized_runtimes: dict[str, int]
+) -> dict[int, str]:
+    line_comment_ast_mapper = CommentMapper(test, original_runtimes, optimized_runtimes)
+    source_code = test.generated_original_test_source
+    tree = ast.parse(source_code)
+    line_comment_ast_mapper.visit(tree)
+    return line_comment_ast_mapper.results
 
-class CfoVisitor(ast.NodeVisitor):
-    """AST visitor that finds all assignments to a variable named 'codeflash_output'.
 
-    and reports their location relative to the function they're in.
-    """
+class CommentAdder(cst.CSTTransformer):
+    """Transformer that adds comments to specified lines."""
 
-    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
+    # Declare metadata dependencies
+    METADATA_DEPENDENCIES = (PositionProvider,)
 
-    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 __init__(self, line_to_comments: dict[int, str]) -> None:
+        """Initialize the transformer with target line numbers.
 
-    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
+        Args:
+            line_to_comments: Mapping of line numbers (1-indexed) to comments
 
+        """
+        self.line_to_comments = line_to_comments
+        super().__init__()
 
-def find_codeflash_output_assignments(function_name: str, source_code: str) -> list[int]:
-    tree = ast.parse(source_code)
-    visitor = CfoVisitor(function_name, source_code)
-    visitor.visit(tree)
-    return visitor.results
+    def leave_SimpleStatementLine(
+        self, original_node: cst.SimpleStatementLine, updated_node: cst.SimpleStatementLine
+    ) -> cst.SimpleStatementLine:
+        """Add comment to simple statement lines."""
+        pos = self.get_metadata(PositionProvider, original_node)
 
+        if pos and pos.start.line in self.line_to_comments:
+            # Create a comment with trailing whitespace
+            comment = cst.TrailingWhitespace(
+                whitespace=cst.SimpleWhitespace(" "), comment=cst.Comment(self.line_to_comments[pos.start.line])
+            )
 
-class Finder(cst.CSTVisitor):
-    def __init__(self, name: str) -> None:
-        super().__init__()
-        self.found = False
-        self.name = name
+            # Update the trailing whitespace of the line itself
+            return updated_node.with_changes(trailing_whitespace=comment)
 
-    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
+        return updated_node
 
+    def leave_SimpleStatementSuite(
+        self, original_node: cst.SimpleStatementSuite, updated_node: cst.SimpleStatementSuite
+    ) -> cst.SimpleStatementSuite:
+        """Add comment to simple statement suites (e.g., after if/for/while)."""
+        pos = self.get_metadata(PositionProvider, original_node)
 
-# 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)
+        if pos and pos.start.line in self.line_to_comments:
+            # Create a comment with trailing whitespace
+            comment = cst.TrailingWhitespace(
+                whitespace=cst.SimpleWhitespace("  "), comment=cst.Comment(self.line_to_comments[pos.start.line])
+            )
 
-    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
+            # Update the trailing whitespace of the suite
+            return updated_node.with_changes(trailing_whitespace=comment)
 
-    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]
-                ):
-                    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
-                        )
-                    )
-                    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})"
-                    )
-                    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
 
-    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
+def unique_inv_id(inv_id_runtimes: dict[InvocationId, list[int]]) -> dict[str, int]:
+    unique_inv_ids: dict[str, int] = {}
+    for inv_id, runtimes in inv_id_runtimes.items():
+        test_qualified_name = (
+            inv_id.test_class_name + "." + inv_id.test_function_name  # type: ignore[operator]
+            if inv_id.test_class_name
+            else inv_id.test_function_name
+        )
+        abs_path = str(Path(inv_id.test_module_path.replace(".", os.sep)).with_suffix(".py").resolve().with_suffix(""))
+        if "__unit_test_" not in abs_path:
+            continue
+        key = test_qualified_name + "#" + abs_path  # type: ignore[operator]
+        parts = inv_id.iteration_id.split("_").__len__()  # type: ignore[union-attr]
+        cur_invid = inv_id.iteration_id.split("_")[0] if parts < 3 else "_".join(inv_id.iteration_id.split("_")[:-1])  # type: ignore[union-attr]
+        match_key = key + "#" + cur_invid
+        if match_key not in unique_inv_ids:
+            unique_inv_ids[match_key] = 0
+        unique_inv_ids[match_key] += min(runtimes)
+    return unique_inv_ids
 
 
 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
-
+    original_runtimes_dict = unique_inv_id(original_runtimes)
+    optimized_runtimes_dict = unique_inv_id(optimized_runtimes)
     # Process each generated test
     modified_tests = []
     for test in generated_tests.generated_tests:
         try:
-            # Parse the test source code
             tree = cst.parse_module(test.generated_original_test_source)
-            # Transform the tree to add runtime comments
-            # 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
+            wrapper = MetadataWrapper(tree)
+            line_to_comments = get_fn_call_linenos(test, original_runtimes_dict, optimized_runtimes_dict)
+            comment_adder = CommentAdder(line_to_comments)
+            modified_tree = wrapper.visit(comment_adder)
             modified_source = modified_tree.code
-
-            # Create a new GeneratedTests object with the modified source
             modified_test = GeneratedTests(
                 generated_original_test_source=modified_source,
                 instrumented_behavior_test_source=test.instrumented_behavior_test_source,
@@ -254,3 +192,28 @@ def add_runtime_comments_to_generated_tests(
             modified_tests.append(test)
 
     return GeneratedTestsList(generated_tests=modified_tests)
+
+
+def remove_functions_from_generated_tests(
+    generated_tests: GeneratedTestsList, test_functions_to_remove: list[str]
+) -> GeneratedTestsList:
+    new_generated_tests = []
+    for generated_test in generated_tests.generated_tests:
+        for test_function in test_functions_to_remove:
+            function_pattern = re.compile(
+                rf"(@pytest\.mark\.parametrize\(.*?\)\s*)?def\s+{re.escape(test_function)}\(.*?\):.*?(?=\ndef\s|$)",
+                re.DOTALL,
+            )
+
+            match = function_pattern.search(generated_test.generated_original_test_source)
+
+            if match is None or "@pytest.mark.parametrize" in match.group(0):
+                continue
+
+            generated_test.generated_original_test_source = function_pattern.sub(
+                "", generated_test.generated_original_test_source
+            )
+
+        new_generated_tests.append(generated_test)
+
+    return GeneratedTestsList(generated_tests=new_generated_tests)
diff --git a/codeflash/code_utils/time_utils.py b/codeflash/code_utils/time_utils.py
index 4e32eedab..e44c279d3 100644
--- a/codeflash/code_utils/time_utils.py
+++ b/codeflash/code_utils/time_utils.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import datetime as dt
 import re
 
@@ -58,42 +60,27 @@ def format_time(nanoseconds: int) -> str:
         raise TypeError("Input must be an integer.")
     if nanoseconds < 0:
         raise ValueError("Input must be a positive integer.")
-    conversions = [(1_000_000_000, "s"), (1_000_000, "ms"), (1_000, "μs"), (1, "ns")]
 
-    # Handle nanoseconds case directly (no decimal formatting needed)
     if nanoseconds < 1_000:
         return f"{nanoseconds}ns"
-
-    # Find appropriate unit
-    for divisor, unit in conversions:
-        if nanoseconds >= divisor:
-            value = nanoseconds / divisor
-            int_value = nanoseconds // divisor
-
-            # Use integer formatting for values >= 100
-            if int_value >= 100:
-                formatted_value = f"{int_value:.0f}"
-            # Format with precision for 3 significant digits
-            elif value >= 100:
-                formatted_value = f"{value:.0f}"
-            elif value >= 10:
-                formatted_value = f"{value:.1f}"
-            else:
-                formatted_value = f"{value:.2f}"
-
-            return f"{formatted_value}{unit}"
-
-    # This should never be reached, but included for completeness
-    return f"{nanoseconds}ns"
+    if nanoseconds < 1_000_000:
+        value = nanoseconds / 1_000
+        return f"{value:.2f}μs" if value < 10 else (f"{value:.1f}μs" if value < 100 else f"{int(value)}μs")
+    if nanoseconds < 1_000_000_000:
+        value = nanoseconds / 1_000_000
+        return f"{value:.2f}ms" if value < 10 else (f"{value:.1f}ms" if value < 100 else f"{int(value)}ms")
+    value = nanoseconds / 1_000_000_000
+    return f"{value:.2f}s" if value < 10 else (f"{value:.1f}s" if value < 100 else f"{int(value)}s")
 
 
 def format_perf(percentage: float) -> str:
     """Format percentage into a human-readable string with 3 significant digits when needed."""
-    percentage_abs = abs(percentage)
-    if percentage_abs >= 100:
+    # Branch order optimized
+    abs_perc = abs(percentage)
+    if abs_perc >= 100:
         return f"{percentage:.0f}"
-    if percentage_abs >= 10:
+    if abs_perc >= 10:
         return f"{percentage:.1f}"
-    if percentage_abs >= 1:
+    if abs_perc >= 1:
         return f"{percentage:.2f}"
     return f"{percentage:.3f}"
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
index 18fbf5064..ef7b215bc 100644
--- a/codeflash/optimization/function_optimizer.py
+++ b/codeflash/optimization/function_optimizer.py
@@ -1016,11 +1016,7 @@ def find_and_process_best_optimization(
                     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(
-                        qualified_name,
-                        self.test_cfg,
-                        generated_tests,
-                        original_runtime_by_test,
-                        optimized_runtime_by_test,
+                        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]
diff --git a/tests/test_add_runtime_comments.py b/tests/test_add_runtime_comments.py
index d2f31e724..da6e49373 100644
--- a/tests/test_add_runtime_comments.py
+++ b/tests/test_add_runtime_comments.py
@@ -34,7 +34,7 @@ def create_test_invocation(
         return FunctionTestInvocation(
             loop_index=loop_index,
             id=InvocationId(
-                test_module_path="tests.test_module",
+                test_module_path="tests.test_module__unit_test_0",
                 test_class_name=None,
                 test_function_name=test_function_name,
                 function_getting_tested="test_function",
@@ -59,12 +59,11 @@ def test_basic_runtime_comment_addition(self, test_config):
     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=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py",
         )
         generated_tests = GeneratedTestsList(generated_tests=[generated_test])
@@ -82,7 +81,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -108,7 +107,7 @@ def helper_function():
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -129,7 +128,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         modified_source = result.generated_tests[0].generated_original_test_source
 
@@ -165,7 +164,7 @@ def test_different_time_formats(self, test_config):
                 generated_original_test_source=test_source,
                 instrumented_behavior_test_source="",
                 instrumented_perf_test_source="",
-                behavior_file_path=test_config.tests_root / "test_module.py",
+                behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
                 perf_file_path=test_config.tests_root / "test_perf.py"
             )
 
@@ -181,9 +180,7 @@ def test_different_time_formats(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(
-                qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes
-            )
+            result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
             modified_source = result.generated_tests[0].generated_original_test_source
             assert f"# {expected_comment}" in modified_source
@@ -201,7 +198,7 @@ def test_missing_test_results(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -215,7 +212,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that no comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -234,7 +231,7 @@ def test_partial_test_results(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -249,7 +246,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that no comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -267,7 +264,7 @@ def test_multiple_runtimes_uses_minimum(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -289,7 +286,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -307,7 +304,7 @@ def test_no_codeflash_output_assignment(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -324,7 +321,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -342,7 +339,7 @@ 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=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
         qualified_name = "bubble_sort"
@@ -359,7 +356,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -385,7 +382,7 @@ def test_multiple_generated_tests(self, test_config):
             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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -393,7 +390,7 @@ def test_multiple_generated_tests(self, test_config):
             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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -413,7 +410,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -432,14 +429,14 @@ def test_preserved_test_attributes(self, test_config):
         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_behavior_path=test_config.tests_root / "test_module__unit_test_0.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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -455,7 +452,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that other attributes are preserved
         modified_test = result.generated_tests[0]
@@ -482,7 +479,7 @@ def test_multistatement_line_handling(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -499,7 +496,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -524,19 +521,18 @@ def test_add_runtime_comments_simple_function(self, test_config):
     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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -546,9 +542,7 @@ def test_add_runtime_comments_simple_function(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''def test_function():
     codeflash_output = some_function() # 1.00s -> 500ms (100% faster)
@@ -571,14 +565,14 @@ def test_function(self):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name="TestClass",
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -589,9 +583,7 @@ def test_function(self):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''class TestClass:
     def test_function(self):
@@ -619,21 +611,21 @@ def test_add_runtime_comments_multiple_assignments(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -643,9 +635,7 @@ def test_add_runtime_comments_multiple_assignments(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''def test_function():
     setup_data = prepare_test()
@@ -672,7 +662,7 @@ def test_add_runtime_comments_no_matching_runtimes(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -690,9 +680,7 @@ def test_add_runtime_comments_no_matching_runtimes(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Source should remain unchanged
         assert len(result.generated_tests) == 1
@@ -714,14 +702,14 @@ def test_add_runtime_comments_no_codeflash_output(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -731,9 +719,7 @@ def test_add_runtime_comments_no_codeflash_output(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Source should remain unchanged
         assert len(result.generated_tests) == 1
@@ -756,7 +742,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config):
             generated_original_test_source=test_source1,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module1.py",
+            behavior_file_path=test_config.tests_root / "test_module1__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf1.py"
         )
 
@@ -764,14 +750,14 @@ def test_add_runtime_comments_multiple_tests(self, test_config):
             generated_original_test_source=test_source2,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module2.py",
+            behavior_file_path=test_config.tests_root / "test_module2__unit_test_0.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_module_path="tests.test_module1__unit_test_0",
             test_class_name=None,
             test_function_name="test_function1",
             function_getting_tested="some_function",
@@ -779,7 +765,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config):
         )
 
         invocation_id2 = InvocationId(
-            test_module_path="tests.test_module2",
+            test_module_path="tests.test_module2__unit_test_0",
             test_class_name=None,
             test_function_name="test_function2",
             function_getting_tested="some_function", # not used in this test throughout the entire test file
@@ -795,9 +781,7 @@ def test_add_runtime_comments_multiple_tests(self, test_config):
             invocation_id2: [800000000],   # 0.8s
         }
 
-        result = add_runtime_comments_to_generated_tests(
-            qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source1 = '''def test_function1():
     codeflash_output = some_function() # 1.00s -> 500ms (100% faster)
@@ -827,14 +811,14 @@ def test_add_runtime_comments_performance_regression(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -842,7 +826,7 @@ def test_add_runtime_comments_performance_regression(self, test_config):
         )
 
         invocation_id2 = InvocationId(
-            test_module_path="tests.test_module",
+            test_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -852,9 +836,7 @@ def test_add_runtime_comments_performance_regression(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''def test_function():
     codeflash_output = some_function() # 1.00s -> 1.50s (33.3% slower)
@@ -880,7 +862,7 @@ def test_basic_runtime_comment_addition_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py",
         )
         generated_tests = GeneratedTestsList(generated_tests=[generated_test])
@@ -898,7 +880,7 @@ def test_basic_runtime_comment_addition_no_cfo(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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -923,7 +905,7 @@ def helper_function():
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -944,7 +926,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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         modified_source = result.generated_tests[0].generated_original_test_source
 
@@ -979,7 +961,7 @@ def test_different_time_formats_no_cfo(self, test_config):
                 generated_original_test_source=test_source,
                 instrumented_behavior_test_source="",
                 instrumented_perf_test_source="",
-                behavior_file_path=test_config.tests_root / "test_module.py",
+                behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
                 perf_file_path=test_config.tests_root / "test_perf.py"
             )
 
@@ -995,9 +977,7 @@ def test_different_time_formats_no_cfo(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(
-                qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes
-            )
+            result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
             modified_source = result.generated_tests[0].generated_original_test_source
             assert f"# {expected_comment}" in modified_source
@@ -1015,7 +995,7 @@ def test_missing_test_results_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1029,7 +1009,7 @@ def test_missing_test_results_no_cfo(self, test_config):
         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)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that no comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -1048,7 +1028,7 @@ def test_partial_test_results_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1063,7 +1043,7 @@ def test_partial_test_results_no_cfo(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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that no comments were added
         modified_source = result.generated_tests[0].generated_original_test_source
@@ -1081,7 +1061,7 @@ def test_multiple_runtimes_uses_minimum_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1103,7 +1083,7 @@ def test_multiple_runtimes_uses_minimum_no_cfo(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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -1121,7 +1101,7 @@ def test_no_codeflash_output_assignment_invalid_iteration_id(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1138,7 +1118,7 @@ def test_no_codeflash_output_assignment_invalid_iteration_id(self, test_config):
         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)
+        result = add_runtime_comments_to_generated_tests(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
@@ -1156,7 +1136,7 @@ def test_invalid_python_code_handling_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
         qualified_name = "bubble_sort"
@@ -1173,7 +1153,7 @@ def test_invalid_python_code_handling_no_cfo(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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(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
@@ -1198,7 +1178,7 @@ def test_multiple_generated_tests_no_cfo(self, test_config):
             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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1206,7 +1186,7 @@ def test_multiple_generated_tests_no_cfo(self, test_config):
             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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1226,7 +1206,7 @@ def test_multiple_generated_tests_no_cfo(self, test_config):
         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)
+        result = add_runtime_comments_to_generated_tests(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
@@ -1245,14 +1225,14 @@ def test_preserved_test_attributes_no_cfo(self, test_config):
         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_behavior_path=test_config.tests_root / "test_module__unit_test_0.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",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1268,7 +1248,7 @@ def test_preserved_test_attributes_no_cfo(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(qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes)
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Check that other attributes are preserved
         modified_test = result.generated_tests[0]
@@ -1295,7 +1275,7 @@ def test_multistatement_line_handling_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1312,7 +1292,7 @@ def test_multistatement_line_handling_no_cfo(self, test_config):
         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)
+        result = add_runtime_comments_to_generated_tests(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
@@ -1341,14 +1321,14 @@ def test_add_runtime_comments_simple_function_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -1358,9 +1338,7 @@ def test_add_runtime_comments_simple_function_no_cfo(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''def test_function():
     result = some_function(); assert result == expected # 1.00s -> 500ms (100% faster)
@@ -1382,14 +1360,14 @@ def test_function(self):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name="TestClass",
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -1400,9 +1378,7 @@ def test_function(self):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''class TestClass:
     def test_function(self):
@@ -1428,21 +1404,21 @@ def test_add_runtime_comments_multiple_assignments_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -1452,9 +1428,7 @@ def test_add_runtime_comments_multiple_assignments_no_cfo(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source = '''def test_function():
     setup_data = prepare_test()
@@ -1479,7 +1453,7 @@ def test_add_runtime_comments_no_matching_runtimes_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf.py"
         )
 
@@ -1487,7 +1461,7 @@ def test_add_runtime_comments_no_matching_runtimes_no_cfo(self, test_config):
 
         # Different invocation ID that won't match
         invocation_id = InvocationId(
-            test_module_path="tests.other_module",
+            test_module_path="tests.other_module__unit_test_0",
             test_class_name=None,
             test_function_name="other_function",
             function_getting_tested="some_other_function",
@@ -1497,9 +1471,7 @@ def test_add_runtime_comments_no_matching_runtimes_no_cfo(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         # Source should remain unchanged
         assert len(result.generated_tests) == 1
@@ -1524,7 +1496,7 @@ def test_add_runtime_comments_multiple_tests_no_cfo(self, test_config):
             generated_original_test_source=test_source1,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module1.py",
+            behavior_file_path=test_config.tests_root / "test_module1__unit_test_0.py",
             perf_file_path=test_config.tests_root / "test_perf1.py"
         )
 
@@ -1532,14 +1504,14 @@ def test_add_runtime_comments_multiple_tests_no_cfo(self, test_config):
             generated_original_test_source=test_source2,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module2.py",
+            behavior_file_path=test_config.tests_root / "test_module2__unit_test_0.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_module_path="tests.test_module1__unit_test_0",
             test_class_name=None,
             test_function_name="test_function1",
             function_getting_tested="some_function",
@@ -1547,7 +1519,7 @@ def test_add_runtime_comments_multiple_tests_no_cfo(self, test_config):
         )
 
         invocation_id2 = InvocationId(
-            test_module_path="tests.test_module2",
+            test_module_path="tests.test_module2__unit_test_0",
             test_class_name=None,
             test_function_name="test_function2",
             function_getting_tested="some_function", # not used in this test throughout the entire test file
@@ -1563,9 +1535,7 @@ def test_add_runtime_comments_multiple_tests_no_cfo(self, test_config):
             invocation_id2: [800000000],   # 0.8s
         }
 
-        result = add_runtime_comments_to_generated_tests(
-            qualified_name, test_config, generated_tests, original_runtimes, optimized_runtimes
-        )
+        result = add_runtime_comments_to_generated_tests(generated_tests, original_runtimes, optimized_runtimes)
 
         expected_source1 = '''def test_function1():
     result = some_function() # 1.00s -> 500ms (100% faster)
@@ -1594,14 +1564,14 @@ def test_add_runtime_comments_performance_regression_no_cfo(self, test_config):
             generated_original_test_source=test_source,
             instrumented_behavior_test_source="",
             instrumented_perf_test_source="",
-            behavior_file_path=test_config.tests_root / "test_module.py",
+            behavior_file_path=test_config.tests_root / "test_module__unit_test_0.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_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -1609,7 +1579,7 @@ def test_add_runtime_comments_performance_regression_no_cfo(self, test_config):
         )
 
         invocation_id2 = InvocationId(
-            test_module_path="tests.test_module",
+            test_module_path="tests.test_module__unit_test_0",
             test_class_name=None,
             test_function_name="test_function",
             function_getting_tested="some_function",
@@ -1619,9 +1589,7 @@ def test_add_runtime_comments_performance_regression_no_cfo(self, test_config):
         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
-        )
+        result = add_runtime_comments_to_generated_tests(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)
@@ -1631,3 +1599,307 @@ def test_add_runtime_comments_performance_regression_no_cfo(self, test_config):
 
         assert len(result.generated_tests) == 1
         assert result.generated_tests[0].generated_original_test_source == expected_source
+
+    def test_runtime_comment_addition_for(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():
+    a = 2
+    for i in range(3):
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2])
+        assert codeflash_output == [1, 2, 3]
+        c = 4
+    d = 5
+"""
+        expected = """def test_bubble_sort():
+    a = 2
+    for i in range(3):
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2]) # 1.80ms -> 1.20ms (50.0% faster)
+        assert codeflash_output == [1, 2, 3]
+        c = 4
+    d = 5
+"""
+        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__unit_test_0.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_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_0')  # 500μs
+        optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='1_2_0')  # 300μs
+        original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='1_2_1')  # 500μs
+        optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='1_2_1')  # 300μs
+        original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='1_2_2')  # 500μs
+        optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_2')  # 300μs
+
+        original_test_results.add(original_invocation1)
+        optimized_test_results.add(optimized_invocation1)
+        original_test_results.add(original_invocation2)
+        optimized_test_results.add(optimized_invocation2)
+        original_test_results.add(original_invocation3)
+        optimized_test_results.add(optimized_invocation3)
+        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(generated_tests, original_runtimes, optimized_runtimes)
+
+        # Check that comments were added
+        modified_source = result.generated_tests[0].generated_original_test_source
+        assert modified_source == expected
+
+    def test_runtime_comment_addition_while(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():
+    i = 0
+    while i<3:
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2])
+        assert codeflash_output == [1, 2, 3]
+        i += 1
+    d = 5
+"""
+        expected = """def test_bubble_sort():
+    i = 0
+    while i<3:
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2]) # 1.80ms -> 1.20ms (50.0% faster)
+        assert codeflash_output == [1, 2, 3]
+        i += 1
+    d = 5
+"""
+        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__unit_test_0.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_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_0')  # 500μs
+        optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='1_2_0')  # 300μs
+        original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='1_2_1')  # 500μs
+        optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='1_2_1')  # 300μs
+        original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='1_2_2')  # 500μs
+        optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_2')  # 300μs
+
+        original_test_results.add(original_invocation1)
+        optimized_test_results.add(optimized_invocation1)
+        original_test_results.add(original_invocation2)
+        optimized_test_results.add(optimized_invocation2)
+        original_test_results.add(original_invocation3)
+        optimized_test_results.add(optimized_invocation3)
+        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(generated_tests, original_runtimes, optimized_runtimes)
+
+        # Check that comments were added
+        modified_source = result.generated_tests[0].generated_original_test_source
+        assert modified_source == expected
+
+    def test_runtime_comment_addition_with(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():
+    i = 0
+    with open('a.txt','rb') as f:
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2])
+        assert codeflash_output == [1, 2, 5]
+        i += 1
+    d = 5
+"""
+        expected = """def test_bubble_sort():
+    i = 0
+    with open('a.txt','rb') as f:
+        b = 3
+        b1 = 6
+        codeflash_output = bubble_sort([3, 1, 2]) # 1.80ms -> 1.20ms (50.0% faster)
+        assert codeflash_output == [1, 2, 5]
+        i += 1
+    d = 5
+"""
+        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__unit_test_0.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_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_0')  # 500μs
+        optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='1_2_0')  # 300μs
+        original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='1_2_1')  # 500μs
+        optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='1_2_1')  # 300μs
+        original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='1_2_2')  # 500μs
+        optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2_2')  # 300μs
+
+        original_test_results.add(original_invocation1)
+        optimized_test_results.add(optimized_invocation1)
+        original_test_results.add(original_invocation2)
+        optimized_test_results.add(optimized_invocation2)
+        original_test_results.add(original_invocation3)
+        optimized_test_results.add(optimized_invocation3)
+        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(generated_tests, original_runtimes, optimized_runtimes)
+
+        # Check that comments were added
+        modified_source = result.generated_tests[0].generated_original_test_source
+        assert modified_source == expected
+
+    def test_runtime_comment_addition_lc(self, test_config):
+        """Test basic functionality of adding runtime comments for list comprehension."""
+        # Create test source code
+        os.chdir(test_config.project_root_path)
+        test_source = """def test_bubble_sort():
+    i = 0
+    codeflash_output = [bubble_sort([3, 1, 2]) for _ in range(3)]
+    assert codeflash_output == [[1,2,3],[1,2,3],[1,2,3]]
+    i += 1
+    d = 5
+"""
+        expected = """def test_bubble_sort():
+    i = 0
+    codeflash_output = [bubble_sort([3, 1, 2]) for _ in range(3)] # 1.80ms -> 1.20ms (50.0% faster)
+    assert codeflash_output == [[1,2,3],[1,2,3],[1,2,3]]
+    i += 1
+    d = 5
+"""
+        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__unit_test_0.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_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_0')  # 500μs
+        optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='1_0')  # 300μs
+        original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='1_1')  # 500μs
+        optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='1_1')  # 300μs
+        original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='1_2')  # 500μs
+        optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2')  # 300μs
+
+        original_test_results.add(original_invocation1)
+        optimized_test_results.add(optimized_invocation1)
+        original_test_results.add(original_invocation2)
+        optimized_test_results.add(optimized_invocation2)
+        original_test_results.add(original_invocation3)
+        optimized_test_results.add(optimized_invocation3)
+        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(generated_tests, original_runtimes, optimized_runtimes)
+
+        # Check that comments were added
+        modified_source = result.generated_tests[0].generated_original_test_source
+        assert modified_source == expected
+
+    def test_runtime_comment_addition_parameterized(self, test_config):
+        """Test basic functionality of adding runtime comments for list comprehension."""
+        # Create test source code
+        os.chdir(test_config.project_root_path)
+        test_source = """@pytest.mark.parametrize(
+    "input, expected_output",
+    [
+        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
+        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
+        (list(reversed(range(50))), list(range(50))),
+    ],
+)
+def test_bubble_sort(input, expected_output):
+    i = 0
+    codeflash_output = bubble_sort(input)
+    assert codeflash_output == expected_output
+    i += 1
+    d = 5
+"""
+        expected = """@pytest.mark.parametrize(
+    "input, expected_output",
+    [
+        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
+        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
+        (list(reversed(range(50))), list(range(50))),
+    ],
+)
+def test_bubble_sort(input, expected_output):
+    i = 0
+    codeflash_output = bubble_sort(input) # 1.80ms -> 1.20ms (50.0% faster)
+    assert codeflash_output == expected_output
+    i += 1
+    d = 5
+"""
+        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__unit_test_0.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_invocation1 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_0')  # 500μs
+        optimized_invocation1 = self.create_test_invocation("test_bubble_sort", 300_000, iteration_id='1_0')  # 300μs
+        original_invocation2 = self.create_test_invocation("test_bubble_sort", 600_000, iteration_id='1_1')  # 500μs
+        optimized_invocation2 = self.create_test_invocation("test_bubble_sort", 400_000, iteration_id='1_1')  # 300μs
+        original_invocation3 = self.create_test_invocation("test_bubble_sort", 700_000, iteration_id='1_2')  # 500μs
+        optimized_invocation3 = self.create_test_invocation("test_bubble_sort", 500_000, iteration_id='1_2')  # 300μs
+
+        original_test_results.add(original_invocation1)
+        optimized_test_results.add(optimized_invocation1)
+        original_test_results.add(original_invocation2)
+        optimized_test_results.add(optimized_invocation2)
+        original_test_results.add(original_invocation3)
+        optimized_test_results.add(optimized_invocation3)
+        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(generated_tests, original_runtimes, optimized_runtimes)
+
+        # Check that comments were added
+        modified_source = result.generated_tests[0].generated_original_test_source
+        assert modified_source == expected
\ No newline at end of file