Rich error output during baseline establishment (CF-646)

codeflash/optimization/function_optimizer.py

@@ -3,6 +3,7 @@
 import ast
 import concurrent.futures
 import os
+import re
 import subprocess
 import time
 import uuid
@@ -1192,6 +1193,12 @@ def run_and_parse_tests(
                 f"stdout: {run_result.stdout}\n"
                 f"stderr: {run_result.stderr}\n"
             )
+            if "ModuleNotFoundError" in run_result.stdout:
+                from rich.text import Text
+
+                match = re.search(r"^.*ModuleNotFoundError.*$", run_result.stdout, re.MULTILINE).group()
+                panel = Panel(Text.from_markup(f"⚠️  {match} ", style="bold red"), expand=False)
+                console.print(panel)
         if testing_type in {TestingMode.BEHAVIOR, TestingMode.PERFORMANCE}:
             results, coverage_results = parse_test_results(
                 test_xml_path=result_file_path,

tests/test_test_runner.py

@@ -1,3 +1,5 @@
+import re
+
 import os
 import tempfile
 from pathlib import Path
@@ -96,3 +98,51 @@ def test_sort():
         )
     assert results[0].did_pass, "Test did not pass as expected"
     result_file.unlink(missing_ok=True)
+
+    code = """import torch
+def sorter(arr):
+    print(torch.ones(1))
+    arr.sort()
+    return arr
+
+def test_sort():
+    arr = [5, 4, 3, 2, 1, 0]
+    output = sorter(arr)
+    assert output == [0, 1, 2, 3, 4, 5]
+    """
+    cur_dir_path = Path(__file__).resolve().parent
+    config = TestConfig(
+        tests_root=cur_dir_path,
+        project_root_path=cur_dir_path,
+        test_framework="pytest",
+        tests_project_rootdir=cur_dir_path.parent,
+    )
+
+    test_env = os.environ.copy()
+    test_env["CODEFLASH_TEST_ITERATION"] = "0"
+    test_env["CODEFLASH_TRACER_DISABLE"] = "1"
+    if "PYTHONPATH" not in test_env:
+        test_env["PYTHONPATH"] = str(config.project_root_path)
+    else:
+        test_env["PYTHONPATH"] += os.pathsep + str(config.project_root_path)
+
+    with tempfile.NamedTemporaryFile(prefix="test_xx", suffix=".py", dir=cur_dir_path) as fp:
+        test_files = TestFiles(
+            test_files=[TestFile(instrumented_behavior_file_path=Path(fp.name), test_type=TestType.EXISTING_UNIT_TEST)]
+        )
+        fp.write(code.encode("utf-8"))
+        fp.flush()
+        result_file, process, _, _ = run_behavioral_tests(
+            test_files,
+            test_framework=config.test_framework,
+            cwd=Path(config.project_root_path),
+            test_env=test_env,
+            pytest_timeout=1,
+            pytest_target_runtime_seconds=1,
+        )
+        results = parse_test_xml(
+            test_xml_file_path=result_file, test_files=test_files, test_config=config, run_result=process
+        )
+    match = re.search(r"^.*ModuleNotFoundError.*$", process.stdout, re.MULTILINE).group()
+    assert match=="E   ModuleNotFoundError: No module named 'torch'"
+    result_file.unlink(missing_ok=True)