Update discover_unit_tests.py

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -10,6 +10,7 @@
 import tempfile
 import unittest
 from collections import defaultdict
+from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Optional
 
@@ -438,17 +439,28 @@ def process_test_files(
 
     function_to_test_map = defaultdict(set)
 
-    with test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
-        progress,
-        task_id,
+    with (
+        ProcessPoolExecutor() as executor,
+        test_files_progress_bar(total=len(file_to_test_map), description="Processing test files") as (
+            progress,
+            task_id,
+        ),
     ):
+        futures_map = {}
         for test_file, functions in file_to_test_map.items():
-            _, local_results = _process_single_test_file(test_file, functions, project_root_path, test_framework)
-
-            for qualified_name, function_called_in_test in local_results:
-                function_to_test_map[qualified_name].add(function_called_in_test)
-
-            progress.advance(task_id)
+            future = executor.submit(_process_single_test_file, test_file, functions, project_root_path, test_framework)
+            futures_map[future] = test_file
+
+        for future in as_completed(futures_map):
+            original_test_file = futures_map[future]
+            try:
+                _, local_results = future.result()
+                for qualified_name, function_called_in_test in local_results:
+                    function_to_test_map[qualified_name].add(function_called_in_test)
+            except Exception as e:
+                logger.error(f"Error processing test file {original_test_file}: {e}")
+            finally:
+                progress.advance(task_id)
 
     function_to_tests_dict = {function: list(tests) for function, tests in function_to_test_map.items()}
     num_discovered_tests = sum(len(tests) for tests in function_to_tests_dict.values())