mirror of
https://github.com/codeflash-ai/codeflash-agent.git
synced 2026-05-04 18:25:19 +00:00
919a673be2
* chore: add gitignore entries for local eval repos, e2e fixtures, and env files * fix: restore clean bubble_sort_method.py test fixture The call-site ID commit re-contaminated this file with instrumentation decorators, causing tests to fail with missing CODEFLASH_LOOP_INDEX. * fix: resolve ruff and mypy errors in codeflash-python - Add import-not-found ignores for optional torch/jax imports - Extract magic column index to _STDOUT_COLUMN_INDEX constant - Fix unused variable in _instrument_sync.py - Cast cpu_time_ns to int for mypy arg-type * fix: add skip markers for optional deps and apply ruff formatting to tests Skip torch/jax/tensorflow tests when those packages are not installed. Move has_module helper to conftest.py for reuse across test files. Apply ruff format to all test files that drifted. * fix: resolve remaining ruff format and mypy errors - Add missing blank line in conftest.py (ruff format) - Remove unused import-untyped ignore on jax import (mypy unused-ignore) - Add type: ignore comments for object-typed SQLite row values * chore: bump codeflash-python to 0.1.1.dev0
1084 lines
36 KiB
Python
1084 lines
36 KiB
Python
from __future__ import annotations
|
|
|
|
import ast
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from codeflash_python._model import (
|
|
FunctionParent,
|
|
FunctionToOptimize,
|
|
TestingMode,
|
|
)
|
|
from codeflash_python.test_discovery.models import CodePosition
|
|
from codeflash_python.testing._instrument_core import FunctionImportedAsVisitor
|
|
from codeflash_python.testing._instrumentation import (
|
|
inject_profiling_into_existing_test,
|
|
)
|
|
|
|
project_root = Path(__file__).parent.resolve()
|
|
|
|
|
|
@pytest.fixture
|
|
def tmp_dir():
|
|
"""Create a temporary directory for test results."""
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
yield Path(tmpdirname)
|
|
|
|
|
|
def _assert_sync_instrumentation_present(source: str) -> None:
|
|
"""Assert that the new sync instrumentation markers are present."""
|
|
assert "_codeflash_call_site.set(" in source
|
|
assert "from codeflash_async_wrapper import _codeflash_call_site" in source
|
|
|
|
|
|
def _assert_old_instrumentation_absent(source: str) -> None:
|
|
"""Assert that old sync instrumentation artifacts are absent."""
|
|
assert "codeflash_wrap" not in source
|
|
assert "codeflash_con" not in source
|
|
assert "codeflash_cur" not in source
|
|
assert "import sqlite3" not in source
|
|
|
|
|
|
def test_perfinjector_bubble_sort(tmp_dir) -> None:
|
|
"""Instrument a unittest bubble sort test with profiling."""
|
|
code = """import unittest
|
|
|
|
from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
def test_sort(self):
|
|
input = [5, 4, 3, 2, 1, 0]
|
|
output = sorter(input)
|
|
self.assertEqual(output, [0, 1, 2, 3, 4, 5])
|
|
|
|
input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
|
|
output = sorter(input)
|
|
self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
|
|
|
|
input = list(reversed(range(5000)))
|
|
self.assertEqual(sorter(input), list(range(5000)))
|
|
"""
|
|
|
|
with (tmp_dir / "test_sort.py").open("w") as f:
|
|
f.write(code)
|
|
f.flush()
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=Path(f.name)
|
|
)
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
Path(f.name),
|
|
[CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
|
|
func,
|
|
Path(f.name).parent,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "sorter" in new_test
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sort" in new_test
|
|
count = new_test.count("_codeflash_call_site.set(")
|
|
assert count == 3
|
|
|
|
|
|
def test_perfinjector_only_replay_test(tmp_dir) -> None:
|
|
"""Instrument a replay test with profiling."""
|
|
code = """import dill as pickle
|
|
import pytest
|
|
from codeflash.tracing.replay_test import get_next_arg_and_return
|
|
from codeflash.validation.equivalence import compare_results
|
|
from packagename.ml.yolo.image_reshaping_utils import prepare_image_for_yolo as packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo
|
|
def test_prepare_image_for_yolo():
|
|
for arg_val_pkl, return_val_pkl in get_next_arg_and_return('/home/saurabh/packagename/traces/first.trace', 3):
|
|
args = pickle.loads(arg_val_pkl)
|
|
return_val_1= pickle.loads(return_val_pkl)
|
|
ret = packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo(**args)
|
|
assert compare_results(return_val_1, ret)
|
|
"""
|
|
|
|
with (tmp_dir / "test_return_values.py").open("w") as f:
|
|
f.write(code)
|
|
f.flush()
|
|
func = FunctionToOptimize(
|
|
function_name="prepare_image_for_yolo",
|
|
parents=(),
|
|
file_path=Path("module.py"),
|
|
)
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
Path(f.name), [CodePosition(10, 14)], func, Path(f.name).parent
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert (
|
|
"packagename_ml_yolo_image_reshaping_utils_prepare_image_for_yolo"
|
|
in new_test
|
|
)
|
|
assert "def test_prepare_image_for_yolo" in new_test
|
|
assert "_codeflash_call_site.set(" in new_test
|
|
|
|
|
|
def test_perfinjector_bubble_sort_results() -> None:
|
|
"""Instrument bubble sort and verify output structure."""
|
|
code = """from code_to_optimize.bubble_sort import sorter
|
|
import datetime
|
|
|
|
|
|
def test_sort():
|
|
input = [5, 4, 3, 2, 1, 0]
|
|
print(datetime.datetime.now().isoformat())
|
|
output = sorter(input)
|
|
assert output == [0, 1, 2, 3, 4, 5]
|
|
|
|
input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
|
|
output = sorter(input)
|
|
assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""
|
|
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
code_path = (
|
|
project_root / "code_to_optimize/bubble_sort.py"
|
|
).resolve()
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(8, 14), CodePosition(12, 14)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 2
|
|
|
|
os.chdir(run_cwd)
|
|
success_perf, new_perf_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(8, 14), CodePosition(12, 14)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success_perf
|
|
assert new_perf_test is not None
|
|
_assert_sync_instrumentation_present(new_perf_test)
|
|
_assert_old_instrumentation_absent(new_perf_test)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_parametrized_results() -> None:
|
|
"""Instrument parametrized bubble sort and verify output structure."""
|
|
code = """from code_to_optimize.bubble_sort import sorter
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input, expected_output",
|
|
[
|
|
([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
|
|
([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
|
|
(list(reversed(range(50))), list(range(50))),
|
|
],
|
|
)
|
|
def test_sort_parametrized(input, expected_output):
|
|
output = sorter(input)
|
|
assert output == expected_output
|
|
"""
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(14, 13)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(14, 13)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "def test_sort_parametrized" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_parametrized_loop_results() -> None:
|
|
"""Instrument parametrized loop bubble sort and verify output structure."""
|
|
code = """from code_to_optimize.bubble_sort import sorter
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input, expected_output",
|
|
[
|
|
([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
|
|
([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
|
|
(list(reversed(range(50))), list(range(50))),
|
|
],
|
|
)
|
|
def test_sort_parametrized_loop(input, expected_output):
|
|
for i in range(2):
|
|
output = sorter(input)
|
|
assert output == expected_output
|
|
"""
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_loop_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(15, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(15, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "def test_sort_parametrized_loop" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_loop_results() -> None:
|
|
"""Instrument loop bubble sort and verify output structure."""
|
|
code = """from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
def test_sort():
|
|
inputs = [[5, 4, 3, 2, 1, 0], [5.0, 4.0, 3.0, 2.0, 1.0, 0.0], list(reversed(range(50)))]
|
|
expected_outputs = [[0, 1, 2, 3, 4, 5], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], list(range(50))]
|
|
|
|
for i in range(3):
|
|
input = inputs[i]
|
|
expected_output = expected_outputs[i]
|
|
output = sorter(input)
|
|
assert output == expected_output"""
|
|
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_loop_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(str(run_cwd))
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(11, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(11, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_unittest_results() -> None:
|
|
"""Instrument unittest bubble sort and verify output structure."""
|
|
code = """import unittest
|
|
|
|
from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
def test_sort(self):
|
|
input = [5, 4, 3, 2, 1, 0]
|
|
output = sorter(input)
|
|
self.assertEqual(output, [0, 1, 2, 3, 4, 5])
|
|
|
|
input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
|
|
output = sorter(input)
|
|
self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])
|
|
|
|
input = list(reversed(range(50)))
|
|
output = sorter(input)
|
|
self.assertEqual(output, list(range(50)))
|
|
"""
|
|
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(9, 17), CodePosition(13, 17), CodePosition(17, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 3
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
assert new_test_perf.count("_codeflash_call_site.set(") == 3
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_unittest_parametrized_results() -> None:
|
|
"""Instrument unittest parametrized bubble sort and verify output structure."""
|
|
code = """import unittest
|
|
from parameterized import parameterized
|
|
|
|
from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
@parameterized.expand(
|
|
[
|
|
([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
|
|
([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
|
|
(list(reversed(range(50))), list(range(50))),
|
|
]
|
|
)
|
|
def test_sort(self, input, expected_output):
|
|
output = sorter(input)
|
|
self.assertEqual(output, expected_output)
|
|
"""
|
|
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
project_root_path = project_root
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(16, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(16, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_unittest_loop_results() -> None:
|
|
"""Instrument unittest loop bubble sort and verify output structure."""
|
|
code = """import unittest
|
|
|
|
from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
def test_sort(self):
|
|
inputs = [[5, 4, 3, 2, 1, 0], [5.0, 4.0, 3.0, 2.0, 1.0, 0.0], list(reversed(range(50)))]
|
|
expected_outputs = [[0, 1, 2, 3, 4, 5], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], list(range(50))]
|
|
|
|
for i in range(3):
|
|
input = inputs[i]
|
|
expected_output = expected_outputs[i]
|
|
output = sorter(input)
|
|
self.assertEqual(output, expected_output)"""
|
|
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_loop_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(14, 21)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
assert success
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(14, 21)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_perfinjector_bubble_sort_unittest_parametrized_loop_results() -> None:
|
|
"""Instrument unittest parametrized loop bubble sort and verify output structure."""
|
|
code = """import unittest
|
|
from parameterized import parameterized
|
|
|
|
from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
@parameterized.expand(
|
|
[
|
|
([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
|
|
([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
|
|
(list(reversed(range(50))), list(range(50))),
|
|
]
|
|
)
|
|
def test_sort(self, input, expected_output):
|
|
for i in range(2):
|
|
output = sorter(input)
|
|
self.assertEqual(output, expected_output)
|
|
"""
|
|
|
|
code_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_loop_results_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
project_root_path = project_root
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="sorter", file_path=code_path, parents=()
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(17, 21)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.BEHAVIOR,
|
|
)
|
|
success_perf, new_test_perf = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(17, 21)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert "sorter" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
assert success_perf
|
|
assert new_test_perf is not None
|
|
_assert_sync_instrumentation_present(new_test_perf)
|
|
_assert_old_instrumentation_absent(new_test_perf)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_class_method_imported_as() -> None:
|
|
"""Detect import aliases for functions and class methods."""
|
|
code = """import functionA
|
|
import moduleB as module_B
|
|
from module import functionB as function_B
|
|
import class_name_B
|
|
from nuitka.nodes.ImportNodes import ExpressionBuiltinImport as nuitka_nodes_ImportNodes_ExpressionBuiltinImport
|
|
"""
|
|
f = FunctionToOptimize(
|
|
function_name="functionA", file_path=Path("module.py"), parents=()
|
|
)
|
|
tree = ast.parse(code)
|
|
visitor = FunctionImportedAsVisitor(f)
|
|
visitor.visit(tree)
|
|
assert visitor.imported_as.function_name == "functionA"
|
|
|
|
f = FunctionToOptimize(
|
|
function_name="functionB", file_path=Path("module.py"), parents=()
|
|
)
|
|
visitor = FunctionImportedAsVisitor(f)
|
|
visitor.visit(tree)
|
|
assert visitor.imported_as.function_name == "function_B"
|
|
|
|
f = FunctionToOptimize(
|
|
function_name="method_name",
|
|
file_path=Path("module.py"),
|
|
parents=(FunctionParent("ExpressionBuiltinImport", "ClassDef"),),
|
|
)
|
|
visitor = FunctionImportedAsVisitor(f)
|
|
visitor.visit(tree)
|
|
assert (
|
|
visitor.imported_as.qualified_name
|
|
== "nuitka_nodes_ImportNodes_ExpressionBuiltinImport.method_name"
|
|
)
|
|
|
|
f = FunctionToOptimize(
|
|
function_name="class_name_B", file_path=Path("module.py"), parents=()
|
|
)
|
|
visitor = FunctionImportedAsVisitor(f)
|
|
visitor.visit(tree)
|
|
assert visitor.imported_as.qualified_name == "class_name_B"
|
|
|
|
|
|
def test_class_function_instrumentation() -> None:
|
|
"""Instrument a class method call in a test function."""
|
|
code = """from module import class_name as class_name_A
|
|
|
|
def test_class_name_A_function_name():
|
|
ret = class_name_A.function_name(**args)
|
|
"""
|
|
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_class_function_instrumentation_temp.py"
|
|
)
|
|
try:
|
|
with open(test_path, "w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root / "code_to_optimize/"
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
func = FunctionToOptimize(
|
|
function_name="function_name",
|
|
file_path=project_root_path / "module.py",
|
|
parents=(FunctionParent("class_name", "ClassDef"),),
|
|
)
|
|
os.chdir(str(run_cwd))
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path, [CodePosition(4, 23)], func, project_root_path
|
|
)
|
|
os.chdir(original_cwd)
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "class_name_A.function_name" in new_test
|
|
assert "def test_class_name_A_function_name" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
|
|
|
|
def test_wrong_function_instrumentation() -> None:
|
|
"""Instrument multiple calls to find_common_tags correctly."""
|
|
code = """from codeflash.result.common_tags import find_common_tags
|
|
|
|
|
|
def test_common_tags_1():
|
|
articles_1 = [1, 2, 3]
|
|
|
|
assert find_common_tags(articles_1) == set(1, 2)
|
|
|
|
articles_2 = [1, 2]
|
|
|
|
assert find_common_tags(articles_2) == set(1)
|
|
"""
|
|
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_wrong_function_instrumentation_temp.py"
|
|
)
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root / "code_to_optimize/"
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
func = FunctionToOptimize(
|
|
function_name="find_common_tags",
|
|
file_path=project_root_path / "module.py",
|
|
parents=(),
|
|
)
|
|
|
|
os.chdir(str(run_cwd))
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(7, 11), CodePosition(11, 11)],
|
|
func,
|
|
project_root_path,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "find_common_tags" in new_test
|
|
assert "def test_common_tags_1" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 2
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_conditional_instrumentation() -> None:
|
|
"""Instrument a function call inside an if block."""
|
|
code = """from code_to_optimize.bubble_sort import sorter
|
|
|
|
|
|
def test_sort():
|
|
input = [5, 4, 3, 2, 1, 0]
|
|
if len(input) > 0:
|
|
assert sorter(input) == [0, 1, 2, 3, 4, 5]"""
|
|
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_conditional_instrumentation_temp.py"
|
|
)
|
|
try:
|
|
with open(test_path, "w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root / "code_to_optimize/"
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
func = FunctionToOptimize(
|
|
function_name="sorter",
|
|
file_path=project_root_path / "module.py",
|
|
parents=(),
|
|
)
|
|
|
|
os.chdir(str(run_cwd))
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path, [CodePosition(7, 15)], func, project_root_path
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "sorter" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_static_method_instrumentation():
|
|
"""Instrument a static method call (BubbleSorter.sorter)."""
|
|
code = """from code_to_optimize.bubble_sort import BubbleSorter
|
|
|
|
|
|
def test_sort():
|
|
input = [5, 4, 3, 2, 1, 0]
|
|
output = BubbleSorter.sorter(input)
|
|
assert output == [0, 1, 2, 3, 4, 5]
|
|
|
|
input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
|
|
output = BubbleSorter.sorter(input)
|
|
assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""
|
|
|
|
function_to_optimize = FunctionToOptimize(
|
|
function_name="sorter",
|
|
file_path=Path(
|
|
"/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py"
|
|
),
|
|
parents=(FunctionParent("BubbleSorter", "ClassDef"),),
|
|
starting_line=None,
|
|
ending_line=None,
|
|
)
|
|
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_temp.py"
|
|
)
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
project_root_path = project_root / "code_to_optimize/"
|
|
run_cwd = project_root
|
|
original_cwd = Path.cwd()
|
|
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(6, 26), CodePosition(10, 26)],
|
|
function_to_optimize,
|
|
project_root_path,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "BubbleSorter.sorter" in new_test
|
|
assert "def test_sort" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 2
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_class_method_instrumentation(tmp_path: Path) -> None:
|
|
"""Instrument a class method call (Optimizer.get_code_optimization_context)."""
|
|
code = """from codeflash.optimization.optimizer import Optimizer
|
|
def test_code_replacement10() -> None:
|
|
get_code_output = '''random code'''
|
|
file_path = Path(__file__).resolve()
|
|
opt = Optimizer(
|
|
Namespace(
|
|
project_root=str(file_path.parent.resolve()),
|
|
disable_telemetry=True,
|
|
tests_root="tests",
|
|
test_framework="pytest",
|
|
pytest_cmd="pytest",
|
|
experiment_id=None,
|
|
),
|
|
)
|
|
func_top_optimize = FunctionToOptimize(
|
|
function_name="main_method",
|
|
file_path=str(file_path),
|
|
parents=[FunctionParent("MainClass", "ClassDef")],
|
|
)
|
|
with open(file_path) as f:
|
|
original_code = f.read()
|
|
code_context = opt.get_code_optimization_context(
|
|
function_to_optimize=func_top_optimize,
|
|
project_root=str(file_path.parent),
|
|
original_source_code=original_code,
|
|
).unwrap()
|
|
assert code_context.testgen_context_code == get_code_output
|
|
code_context = opt.get_code_optimization_context(
|
|
function_to_optimize=func_top_optimize,
|
|
project_root=str(file_path.parent),
|
|
original_source_code=original_code,
|
|
)
|
|
assert code_context.testgen_context_code == get_code_output
|
|
"""
|
|
|
|
test_file_path = tmp_path / "test_class_method_instrumentation.py"
|
|
test_file_path.write_text(code, encoding="utf-8")
|
|
|
|
func = FunctionToOptimize(
|
|
function_name="get_code_optimization_context",
|
|
parents=(FunctionParent("Optimizer", "ClassDef"),),
|
|
file_path=test_file_path,
|
|
)
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_file_path,
|
|
[CodePosition(22, 28), CodePosition(28, 28)],
|
|
func,
|
|
test_file_path.parent,
|
|
)
|
|
os.chdir(original_cwd)
|
|
assert success
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "opt.get_code_optimization_context" in new_test
|
|
assert "def test_code_replacement10" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") >= 1
|
|
|
|
|
|
def test_time_correction_instrumentation() -> None:
|
|
"""Instrument parametrized sleep test for performance timing."""
|
|
code = """from code_to_optimize.sleeptime import accurate_sleepfunc
|
|
import pytest
|
|
@pytest.mark.parametrize("n, expected_total_sleep_time", [
|
|
(0.01, 0.010),
|
|
(0.02, 0.020),
|
|
])
|
|
def test_sleepfunc_sequence_short(n, expected_total_sleep_time):
|
|
output = accurate_sleepfunc(n)
|
|
assert output == expected_total_sleep_time
|
|
|
|
"""
|
|
|
|
code_path = (project_root / "code_to_optimize/sleeptime.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/pytest/test_time_correction_instrumentation_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
func = FunctionToOptimize(
|
|
function_name="accurate_sleepfunc", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(8, 13)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success, "Test instrumentation failed"
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "accurate_sleepfunc" in new_test
|
|
assert "def test_sleepfunc_sequence_short" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|
|
|
|
|
|
def test_time_correction_instrumentation_unittest() -> None:
|
|
"""Instrument parametrized unittest sleep test for performance timing."""
|
|
code = """import unittest
|
|
from parameterized import parameterized
|
|
|
|
from code_to_optimize.sleeptime import accurate_sleepfunc
|
|
|
|
|
|
class TestPigLatin(unittest.TestCase):
|
|
@parameterized.expand([
|
|
(0.01, 0.010),
|
|
(0.02, 0.020),
|
|
])
|
|
def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time):
|
|
output = accurate_sleepfunc(n)
|
|
"""
|
|
|
|
code_path = (project_root / "code_to_optimize/sleeptime.py").resolve()
|
|
test_path = (
|
|
project_root
|
|
/ "code_to_optimize/tests/unittest/test_time_correction_instrumentation_unittest_temp.py"
|
|
).resolve()
|
|
try:
|
|
with test_path.open("w") as f:
|
|
f.write(code)
|
|
|
|
project_root_path = project_root
|
|
original_cwd = Path.cwd()
|
|
run_cwd = project_root
|
|
func = FunctionToOptimize(
|
|
function_name="accurate_sleepfunc", parents=(), file_path=code_path
|
|
)
|
|
os.chdir(run_cwd)
|
|
success, new_test = inject_profiling_into_existing_test(
|
|
test_path,
|
|
[CodePosition(13, 17)],
|
|
func,
|
|
project_root_path,
|
|
mode=TestingMode.PERFORMANCE,
|
|
)
|
|
os.chdir(original_cwd)
|
|
|
|
assert success, "Test instrumentation failed"
|
|
assert new_test is not None
|
|
_assert_sync_instrumentation_present(new_test)
|
|
_assert_old_instrumentation_absent(new_test)
|
|
assert "accurate_sleepfunc" in new_test
|
|
assert "class TestPigLatin" in new_test
|
|
assert "def test_sleepfunc_sequence_short" in new_test
|
|
assert new_test.count("_codeflash_call_site.set(") == 1
|
|
finally:
|
|
test_path.unlink(missing_ok=True)
|