print number of replay tests found

Author: misrasaurabh1

codeflash/discovery/discover_unit_tests.py

@@ -352,8 +352,10 @@ def discover_unit_tests(
     functions_to_optimize = None
     if file_to_funcs_to_optimize:
         functions_to_optimize = [func for funcs_list in file_to_funcs_to_optimize.values() for func in funcs_list]
-    function_to_tests, num_discovered_tests = strategy(cfg, discover_only_these_tests, functions_to_optimize)
-    return function_to_tests, num_discovered_tests
+    function_to_tests, num_discovered_tests, num_discovered_replay_tests = strategy(
+        cfg, discover_only_these_tests, functions_to_optimize
+    )
+    return function_to_tests, num_discovered_tests, num_discovered_replay_tests
 
 
 def discover_tests_pytest(
@@ -515,6 +517,7 @@ def process_test_files(
 
     function_to_test_map = defaultdict(set)
     num_discovered_tests = 0
+    num_discovered_replay_tests = 0
     jedi_project = jedi.Project(path=project_root_path)
 
     with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
@@ -661,11 +664,14 @@ def process_test_files(
                                     position=CodePosition(line_no=name.line, col_no=name.column),
                                 )
                             )
+                            if test_func.test_type == TestType.REPLAY_TEST:
+                                num_discovered_replay_tests += 1
+
                             num_discovered_tests += 1
                 except Exception as e:
                     logger.debug(str(e))
                     continue
 
             progress.advance(task_id)
 
-    return dict(function_to_test_map), num_discovered_tests
+    return dict(function_to_test_map), num_discovered_tests, num_discovered_replay_tests

codeflash/discovery/functions_to_optimize.py

@@ -304,7 +304,7 @@ def get_all_replay_test_functions(
         logger.error("Could not find trace_file_path in replay test files.")
         exit_with_message("Could not find trace_file_path in replay test files.")
 
-    function_tests, _ = discover_unit_tests(test_cfg, discover_only_these_tests=replay_test)
+    function_tests, _, _ = discover_unit_tests(test_cfg, discover_only_these_tests=replay_test)
     # Get the absolute file paths for each function, excluding class name if present
     filtered_valid_functions = defaultdict(list)
     file_to_functions_map = defaultdict(list)

codeflash/optimization/optimizer.py

@@ -226,12 +226,12 @@ def discover_tests(
 
         console.rule()
         start_time = time.time()
-        function_to_tests, num_discovered_tests = discover_unit_tests(
+        function_to_tests, num_discovered_tests, num_discovered_replay_tests = discover_unit_tests(
             self.test_cfg, file_to_funcs_to_optimize=file_to_funcs_to_optimize
         )
         console.rule()
         logger.info(
-            f"Discovered {num_discovered_tests} existing unit tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
+            f"Discovered {num_discovered_tests} existing unit tests and {num_discovered_replay_tests} replay tests in {(time.time() - start_time):.1f}s at {self.test_cfg.tests_root}"
         )
         console.rule()
         ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})