wrote some tests for instrumentation+testrun+parsing, todo, write some tests for each component

Author: aseembits93

tests/test_instrument_tests.py

@@ -12,8 +12,11 @@
     FunctionImportedAsVisitor,
     inject_profiling_into_existing_test,
 )
+from codeflash.code_utils.line_profile_utils import add_decorator_imports
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
-from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestsInFile
+from codeflash.either import is_successful, Failure
+from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestsInFile, \
+    CodeOptimizationContext
 from codeflash.optimization.function_optimizer import FunctionOptimizer
 from codeflash.verification.test_results import TestType
 from codeflash.verification.verification_utils import TestConfig
@@ -488,7 +491,31 @@ def test_sort():
 codeflash stdout: Sorting list
 result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""
         assert out_str == test_results_perf[1].stdout
-
+        ctx_result = func_optimizer.get_code_optimization_context()
+        code_context: CodeOptimizationContext = ctx_result.unwrap()
+        original_helper_code: dict[Path, str] = {}
+        helper_function_paths = {hf.file_path for hf in code_context.helper_functions}
+        for helper_function_path in helper_function_paths:
+            with helper_function_path.open(encoding="utf8") as f:
+                helper_code = f.read()
+                original_helper_code[helper_function_path] = helper_code
+        line_profiler_output_file = add_decorator_imports(
+            func_optimizer.function_to_optimize, code_context)
+        line_profile_results, _ = func_optimizer.run_and_parse_tests(
+            testing_type=TestingMode.LINE_PROFILE,
+            test_env=test_env,
+            test_files=test_files,
+            optimization_iteration=0,
+            pytest_min_loops=1,
+            pytest_max_loops=1,
+            testing_time=0.1,
+            line_profiler_output_file = line_profiler_output_file
+        )
+        func_optimizer.write_code_and_helpers(
+            func_optimizer.function_to_optimize_source_code, original_helper_code, func_optimizer.function_to_optimize.file_path
+        )
+        tmp_lpr = list(line_profile_results['timings'].keys())
+        assert len(tmp_lpr) == 1 and line_profile_results['timings'][tmp_lpr[0]][0][1]==2
     finally:
         test_path.unlink(missing_ok=True)
         test_path_perf.unlink(missing_ok=True)
@@ -714,7 +741,31 @@ def test_sort_parametrized(input, expected_output):
         )
         assert test_results_perf[2].runtime > 0
         assert test_results_perf[2].did_pass
-
+        ctx_result = func_optimizer.get_code_optimization_context()
+        code_context: CodeOptimizationContext = ctx_result.unwrap()
+        original_helper_code: dict[Path, str] = {}
+        helper_function_paths = {hf.file_path for hf in code_context.helper_functions}
+        for helper_function_path in helper_function_paths:
+            with helper_function_path.open(encoding="utf8") as f:
+                helper_code = f.read()
+                original_helper_code[helper_function_path] = helper_code
+        line_profiler_output_file = add_decorator_imports(
+            func_optimizer.function_to_optimize, code_context)
+        line_profile_results, _ = func_optimizer.run_and_parse_tests(
+            testing_type=TestingMode.LINE_PROFILE,
+            test_env=test_env,
+            test_files=test_files,
+            optimization_iteration=0,
+            pytest_min_loops=1,
+            pytest_max_loops=1,
+            testing_time=0.1,
+            line_profiler_output_file = line_profiler_output_file
+        )
+        func_optimizer.write_code_and_helpers(
+            func_optimizer.function_to_optimize_source_code, original_helper_code, func_optimizer.function_to_optimize.file_path
+        )
+        tmp_lpr = list(line_profile_results['timings'].keys())
+        assert len(tmp_lpr) == 1 and line_profile_results['timings'][tmp_lpr[0]][0][1]==3
     finally:
         test_path.unlink(missing_ok=True)
         test_path_perf.unlink(missing_ok=True)
@@ -1026,6 +1077,31 @@ def test_sort_parametrized_loop(input, expected_output):
         )
         assert test_results[5].runtime > 0
         assert test_results[5].did_pass
+        ctx_result = func_optimizer.get_code_optimization_context()
+        code_context: CodeOptimizationContext = ctx_result.unwrap()
+        original_helper_code: dict[Path, str] = {}
+        helper_function_paths = {hf.file_path for hf in code_context.helper_functions}
+        for helper_function_path in helper_function_paths:
+            with helper_function_path.open(encoding="utf8") as f:
+                helper_code = f.read()
+                original_helper_code[helper_function_path] = helper_code
+        line_profiler_output_file = add_decorator_imports(
+            func_optimizer.function_to_optimize, code_context)
+        line_profile_results, _ = func_optimizer.run_and_parse_tests(
+            testing_type=TestingMode.LINE_PROFILE,
+            test_env=test_env,
+            test_files=test_files,
+            optimization_iteration=0,
+            pytest_min_loops=1,
+            pytest_max_loops=1,
+            testing_time=0.1,
+            line_profiler_output_file = line_profiler_output_file
+        )
+        func_optimizer.write_code_and_helpers(
+            func_optimizer.function_to_optimize_source_code, original_helper_code, func_optimizer.function_to_optimize.file_path
+        )
+        tmp_lpr = list(line_profile_results['timings'].keys())
+        assert len(tmp_lpr) == 1 and line_profile_results['timings'][tmp_lpr[0]][0][1]==6
     finally:
         test_path.unlink(missing_ok=True)
         test_path_behavior.unlink(missing_ok=True)
@@ -1271,6 +1347,31 @@ def test_sort():
         )
         assert test_results[2].runtime > 0
         assert test_results[2].did_pass
+        ctx_result = func_optimizer.get_code_optimization_context()
+        code_context: CodeOptimizationContext = ctx_result.unwrap()
+        original_helper_code: dict[Path, str] = {}
+        helper_function_paths = {hf.file_path for hf in code_context.helper_functions}
+        for helper_function_path in helper_function_paths:
+            with helper_function_path.open(encoding="utf8") as f:
+                helper_code = f.read()
+                original_helper_code[helper_function_path] = helper_code
+        line_profiler_output_file = add_decorator_imports(
+            func_optimizer.function_to_optimize, code_context)
+        line_profile_results, _ = func_optimizer.run_and_parse_tests(
+            testing_type=TestingMode.LINE_PROFILE,
+            test_env=test_env,
+            test_files=test_files,
+            optimization_iteration=0,
+            pytest_min_loops=1,
+            pytest_max_loops=1,
+            testing_time=0.1,
+            line_profiler_output_file = line_profiler_output_file
+        )
+        func_optimizer.write_code_and_helpers(
+            func_optimizer.function_to_optimize_source_code, original_helper_code, func_optimizer.function_to_optimize.file_path
+        )
+        tmp_lpr = list(line_profile_results['timings'].keys())
+        assert len(tmp_lpr) == 1 and line_profile_results['timings'][tmp_lpr[0]][0][1]==3
     finally:
         test_path.unlink(missing_ok=True)
         test_path_perf.unlink(missing_ok=True)