first version that does not work correctly

Author: misrasaurabh1

codeflash/optimization/function_optimizer.py

@@ -341,12 +341,6 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
                     optimized_function=best_optimization.candidate.source_code,
                 )
 
-                existing_tests = existing_tests_source_for(
-                    self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root),
-                    function_to_all_tests,
-                    tests_root=self.test_cfg.tests_root,
-                )
-
                 original_code_combined = original_helper_code.copy()
                 original_code_combined[explanation.file_path] = self.function_to_optimize_source_code
                 new_code_combined = new_helper_code.copy()
@@ -369,6 +363,13 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
                     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),
+                        function_to_all_tests,
+                        tests_root=self.test_cfg.tests_root,
+                        original_test_results=original_code_baseline.benchmarking_test_results,
+                        optimized_test_results=best_optimization.winning_benchmarking_test_results,
+                    )
                     if concolic_test_str:
                         generated_tests_str += "\n\n" + concolic_test_str
 

codeflash/result/create_pr.py

@@ -19,20 +19,77 @@
 from codeflash.github.PrComment import FileDiffContent, PrComment
 
 if TYPE_CHECKING:
-    from codeflash.models.models import FunctionCalledInTest
+    from codeflash.models.models import FunctionCalledInTest, TestResults
     from codeflash.result.explanation import Explanation
 
 
 def existing_tests_source_for(
     function_qualified_name_with_modules_from_root: str,
     function_to_tests: dict[str, set[FunctionCalledInTest]],
     tests_root: Path,
+    original_test_results: Optional[TestResults] = None,
+    optimized_test_results: Optional[TestResults] = None,
 ) -> str:
     test_files = function_to_tests.get(function_qualified_name_with_modules_from_root)
     existing_tests_unique = set()
+
     if test_files:
+        # Group test cases by test file
+        test_files_grouped = {}
         for test_file in test_files:
-            existing_tests_unique.add("- " + str(Path(test_file.tests_in_file.test_file).relative_to(tests_root)))
+            file_path = Path(test_file.tests_in_file.test_file)
+            relative_path = str(file_path.relative_to(tests_root))
+
+            if relative_path not in test_files_grouped:
+                test_files_grouped[relative_path] = []
+            test_files_grouped[relative_path].append(test_file)
+
+        # Create detailed report for each test file
+        for relative_path, tests_in_file in sorted(test_files_grouped.items()):
+            file_line = f"- {relative_path}"
+
+            # Add test case details with timing information if available
+            if original_test_results and optimized_test_results:
+                test_case_details = []
+
+                # Use the same pattern as add_runtime_comments_to_generated_tests
+                original_runtime_by_test = original_test_results.usable_runtime_data_by_test_case()
+                optimized_runtime_by_test = optimized_test_results.usable_runtime_data_by_test_case()
+
+                # Collect test function names for this file
+                test_functions_in_file = {test_file.tests_in_file.test_function for test_file in tests_in_file}
+
+                # Create timing report for each test function
+                for test_function_name in sorted(test_functions_in_file):
+                    # Find matching runtime data
+                    original_runtimes = []
+                    optimized_runtimes = []
+
+                    for invocation_id, runtimes in original_runtime_by_test.items():
+                        if invocation_id.test_function_name == test_function_name:
+                            original_runtimes.extend(runtimes)
+
+                    for invocation_id, runtimes in optimized_runtime_by_test.items():
+                        if invocation_id.test_function_name == test_function_name:
+                            optimized_runtimes.extend(runtimes)
+
+                    if original_runtimes and optimized_runtimes:
+                        # Use minimum timing like the generated tests function does
+                        original_time = min(original_runtimes)
+                        optimized_time = min(optimized_runtimes)
+
+                        from codeflash.code_utils.time_utils import format_time
+
+                        original_str = format_time(original_time)
+                        optimized_str = format_time(optimized_time)
+
+                        test_case_details.append(f"    - {test_function_name}: {original_str} -> {optimized_str}")
+
+                if test_case_details:
+                    file_line += "\n" + "\n".join(test_case_details)
+
+            existing_tests_unique.add(file_line)
+
     return "\n".join(sorted(existing_tests_unique))