codeflash-agent/packages/codeflash-python/tests/test_instrument_all_and_run.py

from __future__ import annotations

import importlib
import os
import sys
import tempfile
from pathlib import Path

from codeflash_python._model import (
    FunctionParent,
    FunctionToOptimize,
    TestingMode,
)
from codeflash_python.analysis._formatter import sort_imports
from codeflash_python.test_discovery.models import CodePosition, TestType
from codeflash_python.testing._instrumentation import (
    get_run_tmp_file,
    inject_profiling_into_existing_test,
    instrument_codeflash_capture,
)
from codeflash_python.testing._parse_results import parse_test_results
from codeflash_python.testing._test_runner import run_behavioral_tests
from codeflash_python.testing.models import TestConfig, TestFile, TestFiles
from codeflash_python.verification._verification import compare_test_results

project_root = Path(__file__).parent.resolve()

# Used by cli instrumentation
codeflash_wrap_string = """def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, codeflash_cur, codeflash_con, *args, **kwargs):
    test_id = f'{{codeflash_test_module_name}}:{{codeflash_test_class_name}}:{{codeflash_test_name}}:{{codeflash_line_id}}:{{codeflash_loop_index}}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {{}}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{{codeflash_line_id}}_{{codeflash_test_index}}'
    test_stdout_tag = f"{{codeflash_test_module_name}}:{{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}}{{codeflash_test_name}}:{{codeflash_function_name}}:{{codeflash_loop_index}}:{{invocation_id}}"
    print(f"!$######{{test_stdout_tag}}######$!")
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        cpu_counter = time.thread_time_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
        codeflash_cpu_duration = time.thread_time_ns() - cpu_counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        codeflash_cpu_duration = time.thread_time_ns() - cpu_counter
        exception = e
    gc.enable()
    print(f"!######{{test_stdout_tag}}######!")
    pickled_return_value = pickle.dumps(exception) if exception else pickle.dumps(return_value)
    codeflash_cur.execute('INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)', (codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_loop_index, invocation_id, codeflash_duration, pickled_return_value, 'function_call', codeflash_cpu_duration))
    codeflash_con.commit()
    if exception:
        raise exception
    return return_value
"""


def _run_and_parse(
    test_files: TestFiles,
    test_env: dict[str, str],
    test_config: TestConfig,
) -> list[object]:
    """Run behavioral tests and parse results (replaces Optimizer.run_and_parse_tests)."""
    xml_path, run_result, _, _ = run_behavioral_tests(
        test_files=test_files,
        test_env=test_env,
        cwd=test_config.project_root_path,
        pytest_cmd=test_config.pytest_cmd,
    )
    return parse_test_results(
        test_xml_path=xml_path,
        test_files=test_files,
        test_config=test_config,
        optimization_iteration=0,
        run_result=run_result,
    )


def test_bubble_sort_behavior_results() -> None:
    code = """from code_to_optimize.bubble_sort import sorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    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]"""

    expected = (
        """import gc
import inspect
import os
import sqlite3
import time

import dill as pickle

from code_to_optimize.bubble_sort import sorter


"""
        + codeflash_wrap_string
        + """
def test_sort():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
    codeflash_con = sqlite3.connect(f'{tmp_dir_path}_{{codeflash_iteration}}.sqlite')
    codeflash_cur = codeflash_con.cursor()
    codeflash_cur.execute('CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT, cpu_runtime INTEGER)')
    input = [5, 4, 3, 2, 1, 0]
    _call__bound__arguments = inspect.signature(sorter).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(sorter, '{module_path}', None, 'test_sort', 'sorter', '1', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0, 1, 2, 3, 4, 5]
    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    _call__bound__arguments = inspect.signature(sorter).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(sorter, '{module_path}', None, 'test_sort', 'sorter', '4', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    codeflash_con.close()
"""
    )

    test_path = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_temp.py"
    ).resolve()
    test_path_perf = (
        project_root
        / "code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_results_perf_temp.py"
    ).resolve()
    fto_path = (project_root / "code_to_optimize/bubble_sort.py").resolve()
    original_code = fto_path.read_text("utf-8")
    try:
        with test_path.open("w") as f:
            f.write(code)

        tests_root = (
            project_root / "code_to_optimize/tests/pytest/"
        ).resolve()
        project_root_path = project_root
        original_cwd = Path.cwd()
        run_cwd = project_root
        func = FunctionToOptimize(
            function_name="sorter",
            parents=(),
            file_path=Path(fto_path),
        )
        os.chdir(run_cwd)
        success, new_test = inject_profiling_into_existing_test(
            test_path,
            [CodePosition(6, 13), CodePosition(10, 13)],
            func,
            project_root_path,
            mode=TestingMode.BEHAVIOR,
        )
        os.chdir(original_cwd)
        assert success
        assert new_test is not None
        assert new_test.replace('"', "'") == expected.format(
            module_path="code_to_optimize.tests.pytest.test_perfinjector_bubble_sort_results_temp",
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        ).replace('"', "'")

        with test_path.open("w") as f:
            f.write(new_test)

        # add codeflash capture
        instrument_codeflash_capture(func, {}, tests_root)

        test_env = os.environ.copy()
        test_env["CODEFLASH_TEST_ITERATION"] = "0"
        test_env["CODEFLASH_LOOP_INDEX"] = "1"
        test_type = TestType.EXISTING_UNIT_TEST

        test_config = TestConfig(
            tests_root=tests_root,
            tests_project_rootdir=project_root_path,
            project_root_path=project_root_path,
            test_framework="pytest",
            pytest_cmd="pytest",
        )
        test_files = TestFiles(
            test_files=[
                TestFile(
                    instrumented_behavior_file_path=test_path,
                    test_type=test_type,
                    original_file_path=test_path,
                    benchmarking_file_path=test_path_perf,
                )
            ]
        )
        test_results = _run_and_parse(test_files, test_env, test_config)

        out_str = """codeflash stdout: Sorting list
result: [0, 1, 2, 3, 4, 5]
"""
        assert test_results[0].stdout == out_str
        assert out_str == test_results[0].stdout
        assert test_results[0].id.function_getting_tested == "sorter"
        assert test_results[0].id.iteration_id == "1_0"
        assert test_results[0].id.test_class_name is None
        assert test_results[0].id.test_function_name == "test_sort"
        assert (
            test_results[0].id.test_module_path
            == "code_to_optimize.tests.pytest.test_perfinjector_bubble_sort_results_temp"
        )
        assert test_results[0].runtime > 0
        assert test_results[0].did_pass
        assert test_results[0].return_value == ([0, 1, 2, 3, 4, 5],)
        out_str = """codeflash stdout: Sorting list
result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
"""
        assert out_str == test_results[1].stdout

        assert test_results[1].id.function_getting_tested == "sorter"
        assert test_results[1].id.iteration_id == "4_0"
        assert test_results[1].id.test_class_name is None
        assert test_results[1].id.test_function_name == "test_sort"
        assert (
            test_results[1].id.test_module_path
            == "code_to_optimize.tests.pytest.test_perfinjector_bubble_sort_results_temp"
        )
        assert test_results[1].runtime > 0
        assert test_results[1].did_pass
        out_str = """codeflash stdout: Sorting list
result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
"""
        assert test_results[1].stdout == out_str
        results2 = _run_and_parse(test_files, test_env, test_config)
        out_str = """codeflash stdout: Sorting list
result: [0, 1, 2, 3, 4, 5]
"""
        assert out_str == results2[0].stdout
        match, _ = compare_test_results(test_results, results2)
        assert match
    finally:
        fto_path.write_text(original_code, "utf-8")
        test_path.unlink(missing_ok=True)
        test_path_perf.unlink(missing_ok=True)


def test_method_full_instrumentation() -> None:
    code = """from code_to_optimize.bubble_sort_method import BubbleSorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    sort_class = BubbleSorter()
    output = sort_class.sorter(input)
    assert output == [0, 1, 2, 3, 4, 5]

    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    sort_class = BubbleSorter()
    output = sort_class.sorter(input)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""

    expected = (
        """import gc
import inspect
import os
import sqlite3
import time

import dill as pickle
from code_to_optimize.bubble_sort_method import BubbleSorter


"""
        + codeflash_wrap_string
        + """
def test_sort():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
    codeflash_con = sqlite3.connect(f'{tmp_dir_path}_{{codeflash_iteration}}.sqlite')
    codeflash_cur = codeflash_con.cursor()
    codeflash_cur.execute('CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT, cpu_runtime INTEGER)')
    input = [5, 4, 3, 2, 1, 0]
    sort_class = BubbleSorter()
    _call__bound__arguments = inspect.signature(sort_class.sorter).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(sort_class.sorter, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter', '2', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0, 1, 2, 3, 4, 5]
    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    sort_class = BubbleSorter()
    _call__bound__arguments = inspect.signature(sort_class.sorter).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(sort_class.sorter, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter', '6', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    codeflash_con.close()
"""
    )
    fto_path = (
        project_root / "code_to_optimize/bubble_sort_method.py"
    ).resolve()
    original_code = fto_path.read_text("utf-8")
    fto = FunctionToOptimize(
        function_name="sorter",
        parents=(FunctionParent(name="BubbleSorter", type="ClassDef"),),
        file_path=Path(fto_path),
    )
    with tempfile.TemporaryDirectory() as tmpdirname:
        tmp_test_path = (
            Path(tmpdirname) / "test_class_method_behavior_results_temp.py"
        )
        tmp_test_path.write_text(code, encoding="utf-8")

        success, new_test = inject_profiling_into_existing_test(
            tmp_test_path,
            [CodePosition(7, 13), CodePosition(12, 13)],
            fto,
            tmp_test_path.parent,
        )
    assert success
    assert new_test.replace('"', "'") == sort_imports(
        expected.format(
            module_path=tmp_test_path.stem,
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        ),
        float_to_top=True,
    ).replace('"', "'")
    tests_root = (project_root / "code_to_optimize/tests/pytest/").resolve()
    test_path = tests_root / "test_class_method_behavior_results_temp.py"
    test_path_perf = (
        tests_root / "test_class_method_behavior_results_perf_temp.py"
    )
    project_root_path = project_root

    try:
        new_test = expected.format(
            module_path="code_to_optimize.tests.pytest.test_class_method_behavior_results_temp",
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        )

        with test_path.open("w") as f:
            f.write(new_test)

        # Add codeflash capture
        instrument_codeflash_capture(fto, {}, tests_root)

        test_env = os.environ.copy()
        test_env["CODEFLASH_TEST_ITERATION"] = "0"
        test_env["CODEFLASH_LOOP_INDEX"] = "1"
        test_type = TestType.EXISTING_UNIT_TEST
        test_config = TestConfig(
            tests_root=tests_root,
            tests_project_rootdir=project_root_path,
            project_root_path=project_root_path,
            test_framework="pytest",
            pytest_cmd="pytest",
        )
        test_files = TestFiles(
            test_files=[
                TestFile(
                    instrumented_behavior_file_path=test_path,
                    test_type=test_type,
                    original_file_path=test_path,
                    benchmarking_file_path=test_path_perf,
                )
            ]
        )
        test_results = _run_and_parse(test_files, test_env, test_config)
        assert len(test_results) == 4
        assert (
            test_results[0].id.function_getting_tested
            == "BubbleSorter.__init__"
        )
        assert test_results[0].id.test_function_name == "test_sort"
        assert test_results[0].did_pass
        assert test_results[0].return_value[0] == {"x": 0}
        assert (
            test_results[1].id.function_getting_tested == "BubbleSorter.sorter"
        )
        assert test_results[1].id.iteration_id == "2_0"
        assert test_results[1].id.test_class_name is None
        assert test_results[1].id.test_function_name == "test_sort"
        assert (
            test_results[1].id.test_module_path
            == "code_to_optimize.tests.pytest.test_class_method_behavior_results_temp"
        )
        assert test_results[1].runtime > 0
        assert test_results[1].did_pass
        assert test_results[1].return_value == ([0, 1, 2, 3, 4, 5],)
        out_str = """codeflash stdout : BubbleSorter.sorter() called\n"""
        assert test_results[1].stdout == out_str
        match, _ = compare_test_results(test_results, test_results)
        assert match
        assert (
            test_results[2].id.function_getting_tested
            == "BubbleSorter.__init__"
        )
        assert test_results[2].id.test_function_name == "test_sort"
        assert test_results[2].did_pass
        assert test_results[2].return_value[0] == {"x": 0}

        assert (
            test_results[3].id.function_getting_tested == "BubbleSorter.sorter"
        )
        assert test_results[3].id.iteration_id == "6_0"
        assert test_results[3].id.test_class_name is None
        assert test_results[3].id.test_function_name == "test_sort"
        assert (
            test_results[3].id.test_module_path
            == "code_to_optimize.tests.pytest.test_class_method_behavior_results_temp"
        )
        assert test_results[3].runtime > 0
        assert test_results[3].did_pass
        assert (
            test_results[3].stdout
            == """codeflash stdout : BubbleSorter.sorter() called\n"""
        )

        results2 = _run_and_parse(test_files, test_env, test_config)

        match, _ = compare_test_results(test_results, results2)
        assert match

        # Replace with optimized code that mutated instance attribute
        optimized_code = """
class BubbleSorter:
    def __init__(self, x=1):
        self.x = x

    def sorter(self, arr):
        for i in range(len(arr)):
            for j in range(len(arr) - 1):
                if arr[j] > arr[j + 1]:
                    temp = arr[j]
                    arr[j] = arr[j + 1]
                    arr[j + 1] = temp
        return arr

        """

        fto_path.write_text(optimized_code, "utf-8")

        # Force reload of module
        module_name = "code_to_optimize.bubble_sort_method"
        if module_name not in sys.modules:
            __import__(module_name)
        importlib.reload(sys.modules[module_name])

        # Add codeflash capture
        instrument_codeflash_capture(fto, {}, tests_root)
        test_config = TestConfig(
            tests_root=tests_root,
            tests_project_rootdir=project_root_path,
            project_root_path=project_root_path,
            test_framework="pytest",
            pytest_cmd="pytest",
        )
        test_files = TestFiles(
            test_files=[
                TestFile(
                    instrumented_behavior_file_path=test_path,
                    test_type=test_type,
                    original_file_path=test_path,
                    benchmarking_file_path=test_path_perf,
                )
            ]
        )
        new_test_results = _run_and_parse(test_files, test_env, test_config)
        assert len(new_test_results) == 4
        assert (
            new_test_results[0].id.function_getting_tested
            == "BubbleSorter.__init__"
        )
        assert new_test_results[0].id.test_function_name == "test_sort"
        assert new_test_results[0].did_pass
        assert new_test_results[0].return_value[0] == {"x": 1}

        assert (
            new_test_results[1].id.function_getting_tested
            == "BubbleSorter.sorter"
        )
        assert new_test_results[1].id.iteration_id == "2_0"
        assert new_test_results[1].id.test_class_name is None
        assert new_test_results[1].id.test_function_name == "test_sort"
        assert (
            new_test_results[1].id.test_module_path
            == "code_to_optimize.tests.pytest.test_class_method_behavior_results_temp"
        )
        assert new_test_results[1].runtime > 0
        assert new_test_results[1].did_pass
        assert new_test_results[1].return_value == ([0, 1, 2, 3, 4, 5],)

        assert (
            new_test_results[2].id.function_getting_tested
            == "BubbleSorter.__init__"
        )
        assert new_test_results[2].id.test_function_name == "test_sort"
        assert new_test_results[2].did_pass
        assert new_test_results[2].return_value[0] == {"x": 1}

        assert (
            new_test_results[3].id.function_getting_tested
            == "BubbleSorter.sorter"
        )
        assert new_test_results[3].id.iteration_id == "6_0"
        assert new_test_results[3].id.test_class_name is None
        assert new_test_results[3].id.test_function_name == "test_sort"
        assert (
            new_test_results[3].id.test_module_path
            == "code_to_optimize.tests.pytest.test_class_method_behavior_results_temp"
        )
        assert new_test_results[3].runtime > 0
        assert new_test_results[3].did_pass
        match, _ = compare_test_results(test_results, new_test_results)
        assert not match

    finally:
        fto_path.write_text(original_code, "utf-8")
        test_path.unlink(missing_ok=True)
        test_path_perf.unlink(missing_ok=True)


def test_classmethod_full_instrumentation() -> None:
    code = """from code_to_optimize.bubble_sort_method import BubbleSorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    output = BubbleSorter.sorter_classmethod(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_classmethod(input)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""

    expected = (
        """import gc
import inspect
import os
import sqlite3
import time

import dill as pickle
from code_to_optimize.bubble_sort_method import BubbleSorter


"""
        + codeflash_wrap_string
        + """
def test_sort():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
    codeflash_con = sqlite3.connect(f'{tmp_dir_path}_{{codeflash_iteration}}.sqlite')
    codeflash_cur = codeflash_con.cursor()
    codeflash_cur.execute('CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT, cpu_runtime INTEGER)')
    input = [5, 4, 3, 2, 1, 0]
    _call__bound__arguments = inspect.signature(BubbleSorter.sorter_classmethod).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(BubbleSorter.sorter_classmethod, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter_classmethod', '1', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0, 1, 2, 3, 4, 5]
    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    _call__bound__arguments = inspect.signature(BubbleSorter.sorter_classmethod).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(BubbleSorter.sorter_classmethod, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter_classmethod', '4', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    codeflash_con.close()
"""
    )
    fto_path = (
        project_root / "code_to_optimize/bubble_sort_method.py"
    ).resolve()
    original_code = fto_path.read_text("utf-8")
    fto = FunctionToOptimize(
        function_name="sorter_classmethod",
        parents=(FunctionParent(name="BubbleSorter", type="ClassDef"),),
        file_path=Path(fto_path),
    )
    with tempfile.TemporaryDirectory() as tmpdirname:
        tmp_test_path = (
            Path(tmpdirname) / "test_classmethod_behavior_results_temp.py"
        )
        tmp_test_path.write_text(code, encoding="utf-8")

        success, new_test = inject_profiling_into_existing_test(
            tmp_test_path,
            [CodePosition(6, 13), CodePosition(10, 13)],
            fto,
            tmp_test_path.parent,
        )
    assert success
    assert new_test.replace('"', "'") == sort_imports(
        expected.format(
            module_path=tmp_test_path.stem,
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        ),
        float_to_top=True,
    ).replace('"', "'")
    tests_root = (project_root / "code_to_optimize/tests/pytest/").resolve()
    test_path = tests_root / "test_classmethod_behavior_results_temp.py"
    test_path_perf = (
        tests_root / "test_classmethod_behavior_results_perf_temp.py"
    )
    project_root_path = project_root

    try:
        new_test = expected.format(
            module_path="code_to_optimize.tests.pytest.test_classmethod_behavior_results_temp",
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        )

        with test_path.open("w") as f:
            f.write(new_test)

        # Add codeflash capture
        instrument_codeflash_capture(fto, {}, tests_root)

        test_env = os.environ.copy()
        test_env["CODEFLASH_TEST_ITERATION"] = "0"
        test_env["CODEFLASH_LOOP_INDEX"] = "1"
        test_type = TestType.EXISTING_UNIT_TEST
        test_config = TestConfig(
            tests_root=tests_root,
            tests_project_rootdir=project_root_path,
            project_root_path=project_root_path,
            test_framework="pytest",
            pytest_cmd="pytest",
        )
        test_files = TestFiles(
            test_files=[
                TestFile(
                    instrumented_behavior_file_path=test_path,
                    test_type=test_type,
                    original_file_path=test_path,
                    benchmarking_file_path=test_path_perf,
                )
            ]
        )
        test_results = _run_and_parse(test_files, test_env, test_config)
        assert len(test_results) == 2
        assert (
            test_results[0].id.function_getting_tested
            == "BubbleSorter.sorter_classmethod"
        )
        assert test_results[0].id.iteration_id == "1_0"
        assert test_results[0].id.test_class_name is None
        assert test_results[0].id.test_function_name == "test_sort"
        assert (
            test_results[0].id.test_module_path
            == "code_to_optimize.tests.pytest.test_classmethod_behavior_results_temp"
        )
        assert test_results[0].runtime > 0
        assert test_results[0].did_pass
        assert test_results[0].return_value == ([0, 1, 2, 3, 4, 5],)
        out_str = """codeflash stdout : BubbleSorter.sorter_classmethod() called
"""
        assert test_results[0].stdout == out_str
        match, _ = compare_test_results(test_results, test_results)
        assert match

        assert (
            test_results[1].id.function_getting_tested
            == "BubbleSorter.sorter_classmethod"
        )
        assert test_results[1].id.iteration_id == "4_0"
        assert test_results[1].id.test_class_name is None
        assert test_results[1].id.test_function_name == "test_sort"
        assert (
            test_results[1].id.test_module_path
            == "code_to_optimize.tests.pytest.test_classmethod_behavior_results_temp"
        )
        assert test_results[1].runtime > 0
        assert test_results[1].did_pass
        assert (
            test_results[1].stdout
            == """codeflash stdout : BubbleSorter.sorter_classmethod() called
"""
        )

        results2 = _run_and_parse(test_files, test_env, test_config)

        match, _ = compare_test_results(test_results, results2)
        assert match

    finally:
        fto_path.write_text(original_code, "utf-8")
        test_path.unlink(missing_ok=True)
        test_path_perf.unlink(missing_ok=True)


def test_staticmethod_full_instrumentation() -> None:
    code = """from code_to_optimize.bubble_sort_method import BubbleSorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    output = BubbleSorter.sorter_staticmethod(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_staticmethod(input)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""

    expected = (
        """import gc
import inspect
import os
import sqlite3
import time

import dill as pickle
from code_to_optimize.bubble_sort_method import BubbleSorter


"""
        + codeflash_wrap_string
        + """
def test_sort():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
    codeflash_con = sqlite3.connect(f'{tmp_dir_path}_{{codeflash_iteration}}.sqlite')
    codeflash_cur = codeflash_con.cursor()
    codeflash_cur.execute('CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT, cpu_runtime INTEGER)')
    input = [5, 4, 3, 2, 1, 0]
    _call__bound__arguments = inspect.signature(BubbleSorter.sorter_staticmethod).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(BubbleSorter.sorter_staticmethod, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter_staticmethod', '1', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0, 1, 2, 3, 4, 5]
    input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
    _call__bound__arguments = inspect.signature(BubbleSorter.sorter_staticmethod).bind(input)
    _call__bound__arguments.apply_defaults()
    output = codeflash_wrap(BubbleSorter.sorter_staticmethod, '{module_path}', None, 'test_sort', 'BubbleSorter.sorter_staticmethod', '4', codeflash_loop_index, codeflash_cur, codeflash_con, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    assert output == [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    codeflash_con.close()
"""
    )
    fto_path = (
        project_root / "code_to_optimize/bubble_sort_method.py"
    ).resolve()
    original_code = fto_path.read_text("utf-8")
    fto = FunctionToOptimize(
        function_name="sorter_staticmethod",
        parents=(FunctionParent(name="BubbleSorter", type="ClassDef"),),
        file_path=Path(fto_path),
    )
    with tempfile.TemporaryDirectory() as tmpdirname:
        tmp_test_path = (
            Path(tmpdirname) / "test_staticmethod_behavior_results_temp.py"
        )
        tmp_test_path.write_text(code, encoding="utf-8")

        success, new_test = inject_profiling_into_existing_test(
            tmp_test_path,
            [CodePosition(6, 13), CodePosition(10, 13)],
            fto,
            tmp_test_path.parent,
        )
    assert success
    assert new_test.replace('"', "'") == sort_imports(
        expected.format(
            module_path=tmp_test_path.stem,
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        ),
        float_to_top=True,
    ).replace('"', "'")
    tests_root = (project_root / "code_to_optimize/tests/pytest/").resolve()
    test_path = tests_root / "test_staticmethod_behavior_results_temp.py"
    test_path_perf = (
        tests_root / "test_staticmethod_behavior_results_perf_temp.py"
    )
    project_root_path = project_root

    try:
        new_test = expected.format(
            module_path="code_to_optimize.tests.pytest.test_staticmethod_behavior_results_temp",
            tmp_dir_path=get_run_tmp_file(
                Path("test_return_values")
            ).as_posix(),
        )

        with test_path.open("w") as f:
            f.write(new_test)

        # Add codeflash capture
        instrument_codeflash_capture(fto, {}, tests_root)

        test_env = os.environ.copy()
        test_env["CODEFLASH_TEST_ITERATION"] = "0"
        test_env["CODEFLASH_LOOP_INDEX"] = "1"
        test_type = TestType.EXISTING_UNIT_TEST
        test_config = TestConfig(
            tests_root=tests_root,
            tests_project_rootdir=project_root_path,
            project_root_path=project_root_path,
            test_framework="pytest",
            pytest_cmd="pytest",
        )
        test_files = TestFiles(
            test_files=[
                TestFile(
                    instrumented_behavior_file_path=test_path,
                    test_type=test_type,
                    original_file_path=test_path,
                    benchmarking_file_path=test_path_perf,
                )
            ]
        )
        test_results = _run_and_parse(test_files, test_env, test_config)
        assert len(test_results) == 2
        assert (
            test_results[0].id.function_getting_tested
            == "BubbleSorter.sorter_staticmethod"
        )
        assert test_results[0].id.iteration_id == "1_0"
        assert test_results[0].id.test_class_name is None
        assert test_results[0].id.test_function_name == "test_sort"
        assert (
            test_results[0].id.test_module_path
            == "code_to_optimize.tests.pytest.test_staticmethod_behavior_results_temp"
        )
        assert test_results[0].runtime > 0
        assert test_results[0].did_pass
        assert test_results[0].return_value == ([0, 1, 2, 3, 4, 5],)
        out_str = """codeflash stdout : BubbleSorter.sorter_staticmethod() called
"""
        assert test_results[0].stdout == out_str
        match, _ = compare_test_results(test_results, test_results)
        assert match

        assert (
            test_results[1].id.function_getting_tested
            == "BubbleSorter.sorter_staticmethod"
        )
        assert test_results[1].id.iteration_id == "4_0"
        assert test_results[1].id.test_class_name is None
        assert test_results[1].id.test_function_name == "test_sort"
        assert (
            test_results[1].id.test_module_path
            == "code_to_optimize.tests.pytest.test_staticmethod_behavior_results_temp"
        )
        assert test_results[1].runtime > 0
        assert test_results[1].did_pass
        assert (
            test_results[1].stdout
            == """codeflash stdout : BubbleSorter.sorter_staticmethod() called
"""
        )

        results2 = _run_and_parse(test_files, test_env, test_config)

        match, _ = compare_test_results(test_results, results2)
        assert match

    finally:
        fto_path.write_text(original_code, "utf-8")
        test_path.unlink(missing_ok=True)
        test_path_perf.unlink(missing_ok=True)