from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize

from testgen.instrumentation.edit_generated_test import parse_module_to_cst
from testgen.instrumentation.instrument_new_tests import instrument_test_source
from testgen.postprocessing.removeassert_transformer import remove_asserts_from_test


def test_instrument_test_source():
    test_source = """import pytest
from code_to_optimize.bubble_sort import sorter

def test_single_element_list():
    # Test that a single element list returns the same single element
    codeflash_output = sorter([42])

def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(1000, 0, -1))  # List from 1000 to 1
    codeflash_output = sorter(large_list)
"""
    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import pytest

def test_single_element_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sorter).bind([42])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_single_element_list', 'sorter', '1', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_single_element_list:sorter:1', codeflash_loop_index, codeflash_invocation_id)

def test_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_list = list(range(1000, 0, -1))
    _call__bound__arguments = inspect.signature(sorter).bind(large_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list', 'sorter', '2', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list:sorter:2', codeflash_loop_index, codeflash_invocation_id)"""
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert result == expected_result


def test_instrument_test_source2():
    test_source = """import pytest
from code_to_optimize.bubble_sort_in_class import BubbleSortClass
bubble_sort_instance = BubbleSortClass()

def test_single_element_list():
    codeflash_output = bubble_sort_instance.sorter([42])

def test_large_list():
    large_list = list(range(1000, 0, -1))
    already_sorted_large_list = list(range(1, 1001))
    codeflash_output = bubble_sort_instance.sorter(large_list)
"""
    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort_in_class.py",
        parents=[FunctionParent(name="BubbleSortClass", type="ClassDef")],
        starting_line=None,
        ending_line=None,
    )
    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort_in_class import BubbleSortClass
import pytest
bubble_sort_instance = BubbleSortClass()

def test_single_element_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(bubble_sort_instance.sorter).bind([42])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(bubble_sort_instance.sorter, 'tests.test_sorter__unit_test_0', None, 'test_single_element_list', 'sorter', '1', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'tests.test_sorter__unit_test_0:test_single_element_list:sorter:1', codeflash_loop_index, codeflash_invocation_id)

def test_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_list = list(range(1000, 0, -1))
    already_sorted_large_list = list(range(1, 1001))
    _call__bound__arguments = inspect.signature(bubble_sort_instance.sorter).bind(large_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(bubble_sort_instance.sorter, 'tests.test_sorter__unit_test_0', None, 'test_large_list', 'sorter', '3', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'tests.test_sorter__unit_test_0:test_large_list:sorter:3', codeflash_loop_index, codeflash_invocation_id)"""

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort_in_class",
        test_module_path="tests.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert result == expected_result


def test_instrument_random_test():
    test_source = """import random
def test_large_list_with_random_values():
    random_list = random.sample(range(1000), 50)
    sorted_list = sorted(random_list)
    codeflash_output = sorter(random_list)"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random

def test_large_list_with_random_values():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    random_list = random.sample(range(1000), 50)
    sorted_list = sorted(random_list)
    _call__bound__arguments = inspect.signature(sorter).bind(random_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list_with_random_values', 'sorter', '4', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list_with_random_values:sorter:4', codeflash_loop_index, codeflash_invocation_id)"""
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert result == expected_result

    test_source = """from random import sample
def test_large_list_with_random_values():
    random_list = sample(range(1000), 50)
    sorted_list = sorted(random_list)
    codeflash_output = sorter(random_list)"""

    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random
from random import sample

def test_large_list_with_random_values():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    random_list = sample(range(1000), 50)
    sorted_list = sorted(random_list)
    _call__bound__arguments = inspect.signature(sorter).bind(random_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list_with_random_values', 'sorter', '4', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list_with_random_values:sorter:4', codeflash_loop_index, codeflash_invocation_id)"""
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert result == expected_result

    test_source = """import pytest
def test_performance_random_large_list():
    import random
    random_list = random.sample(range(1000000), 1000000)

def test_performance_random_large_list2():
    from random import sample
    random_list = sample(range(1000000), 1000000)
"""
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random
import pytest

def test_performance_random_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    random_list = random.sample(range(1000000), 1000000)

def test_performance_random_large_list2():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    from random import sample
    random.seed(42)
    random_list = sample(range(1000000), 1000000)"""
    )

    test_source = """def test_random_from_import_inside_function():
    from random import randint
    value = randint(1, 100)
    assert 1 <= value <= 100
"""
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random

def test_random_from_import_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    from random import randint
    random.seed(42)
    value = randint(1, 100)
    assert 1 <= value <= 100"""
    )

    test_source = """def test_random_alias_import_inside_function():
        import random as rnd
        value = rnd.uniform(0, 1)
        assert 0 <= value <= 1
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_alias_import_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random as rnd
    rnd.seed(42)
    value = rnd.uniform(0, 1)
    assert 0 <= value <= 1"""
    )

    test_source = """def test_random_from_import_with_alias_inside_function():
        from random import random as rnd_func
        value = rnd_func()
        assert 0 <= value <= 1
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random

def test_random_from_import_with_alias_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    from random import random as rnd_func
    random.seed(42)
    value = rnd_func()
    assert 0 <= value <= 1"""
    )

    test_source = """def test_numpy_random_import_inside_function():
        import numpy.random
        value = numpy.random.rand()
        assert 0 <= value <= 1
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_numpy_random_import_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import numpy.random
    numpy.random.seed(42)
    value = numpy.random.rand()
    assert 0 <= value <= 1"""
    )

    test_source = """def test_numpy_random_from_import_inside_function():
        from numpy.random import rand
        value = rand()
        assert 0 <= value <= 1
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import numpy.random

def test_numpy_random_from_import_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    from numpy.random import rand
    numpy.random.seed(42)
    value = rand()
    assert 0 <= value <= 1"""
    )

    test_source = """def test_numpy_random_alias_import_inside_function():
        import numpy.random as nprnd
        value = nprnd.randint(0, 5)
        assert 0 <= value < 5
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_numpy_random_alias_import_inside_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import numpy.random as nprnd
    nprnd.seed(42)
    value = nprnd.randint(0, 5)
    assert 0 <= value < 5"""
    )

    test_source = """def test_random_in_loop():
        import random
        results = []
        for _ in range(5):
            results.append(random.choice(['a', 'b', 'c']))
        assert len(results) == 5
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_in_loop():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    results = []
    for _ in range(5):
        results.append(random.choice(['a', 'b', 'c']))
    assert len(results) == 5"""
    )

    test_source = """def test_random_in_conditional():
        import random
        if random.random() > 0.5:
            outcome = "Heads"
        else:
            outcome = "Tails"
        assert outcome in ["Heads", "Tails"]
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_in_conditional():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    if random.random() > 0.5:
        outcome = 'Heads'
    else:
        outcome = 'Tails'
    assert outcome in ['Heads', 'Tails']"""
    )

    test_source = """def test_random_in_try_except():
        import random
        try:
            value = random.gauss(0, 1)
        except Exception as e:
            value = None
        assert value is not None
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_in_try_except():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    try:
        value = random.gauss(0, 1)
    except Exception as e:
        value = None
    assert value is not None"""
    )

    test_source = """def test_random_with_different_seed():
        import random
        random.seed(100)
        value = random.randint(1, 10)
        assert value == 10 # Known output with seed 100
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_with_different_seed():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    random.seed(100)
    value = random.randint(1, 10)
    assert value == 10"""
    )

    test_source = """import random
def test_random_import_outside_function():
    value = random.randint(1, 10)
    assert 1 <= value <= 10
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random

def test_random_import_outside_function():
    random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    value = random.randint(1, 10)
    assert 1 <= value <= 10"""
    )

    test_source = """def test_random_multiple_imports():
        import random
        import numpy.random as npr
        value1 = random.randint(0, 10)
        value2 = npr.randint(0, 10)
        assert 0 <= value1 <= 10
        assert 0 <= value2 < 10
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_multiple_imports():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    import numpy.random as npr
    npr.seed(42)
    value1 = random.randint(0, 10)
    value2 = npr.randint(0, 10)
    assert 0 <= value1 <= 10
    assert 0 <= value2 < 10"""
    )

    test_source = """def test_random_import_from_and_normal():
        from random import randint
        import random
        value1 = randint(1, 5)
        value2 = random.randint(6, 10)
        assert 1 <= value1 <= 5
        assert 6 <= value2 <= 10
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import random

def test_random_import_from_and_normal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    from random import randint
    random.seed(42)
    import random
    random.seed(42)
    value1 = randint(1, 5)
    value2 = random.randint(6, 10)
    assert 1 <= value1 <= 5
    assert 6 <= value2 <= 10"""
    )

    test_source = """def test_random_with_aliases_and_nested_calls():
        import random as rnd
        def nested():
            return rnd.choice(['x', 'y', 'z'])
        value = nested()
        assert value in ['x', 'y', 'z']
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_with_aliases_and_nested_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random as rnd
    rnd.seed(42)

    def nested():
        return rnd.choice(['x', 'y', 'z'])
    value = nested()
    assert value in ['x', 'y', 'z']"""
    )

    test_source = """def test_random_in_with_statement():
        import random
        with open('dummy_file.txt', 'w') as f:
            f.write(str(random.random()))
        assert True  # Dummy assertion
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_in_with_statement():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    with open('dummy_file.txt', 'w') as f:
        f.write(str(random.random()))
    assert True"""
    )

    test_source = """def test_random_in_try_except_finally():
        import random
        try:
            value = random.expovariate(1)
        except Exception:
            value = None
        finally:
            cleanup_done = True
        assert value is not None and cleanup_done
    """
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert (
        result
        == """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_random_in_try_except_finally():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    import random
    random.seed(42)
    try:
        value = random.expovariate(1)
    except Exception:
        value = None
    finally:
        cleanup_done = True
    assert value is not None and cleanup_done"""
    )
    # TODO: Capture these test cases in future by Recursively process nested functions by visiting FunctionDef & AsyncFunctionDef nodes within the function body
    test_source = """import asyncio
async def test_random_in_async_function():
    import random
    await asyncio.sleep(0.1)
    value = random.randint(1, 5)
    assert 1 <= value <= 5
    """
    test_source = """def test_random_in_nested_function():
        def inner_function():
            import random
            return random.randint(1, 5)
        value = inner_function()
        assert 1 <= value <= 5
    """


def test_instrument_parameterized_test():
    test_source = """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):
    codeflash_output = sorter(input)
"""
    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )
    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
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):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sorter).bind(input)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_sort_parametrized', 'sorter', '1', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_sort_parametrized:sorter:1', codeflash_loop_index, codeflash_invocation_id)"""

    assert result == expected_result


def test_instrument_parameterized_test_with_loop():
    test_source = """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_loop_parametrized(input, expected_output):
    for i in range(2):
        codeflash_output = sorter(input)"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
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_loop_parametrized(input, expected_output):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    for i in range(2):
        _call__bound__arguments = inspect.signature(sorter).bind(input)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_sort_loop_parametrized', 'sorter', '1_0', codeflash_loop_index, **_call__bound__arguments.arguments)
        _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_sort_loop_parametrized:sorter:1_0', codeflash_loop_index, codeflash_invocation_id)"""

    assert result == expected_result


def test_conditional_instrumentation():
    test_source = """from code_to_optimize.bubble_sort import sorter


def test_sort():
    input = [5, 4, 3, 2, 1, 0]
    if len(input) > 0:
        codeflash_output = sorter(input)"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    expected_result = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter

def test_sort():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    input = [5, 4, 3, 2, 1, 0]
    if len(input) > 0:
        _call__bound__arguments = inspect.signature(sorter).bind(input)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_sort', 'sorter', '2_0', codeflash_loop_index, **_call__bound__arguments.arguments)
        _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_sort:sorter:2_0', codeflash_loop_index, codeflash_invocation_id)"""

    assert result == expected_result


def test_remove_asserts_and_instrument_test_source():
    test_source = """import pytest
from code_to_optimize.bubble_sort import sorter

def test_single_element_list():
    # Test that a single element list returns the same single element
    assert sorter([42]) == [42]

def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(1000, 0, -1))  # List from 1000 to 1
    assert sorter(large_list) == sorted(large_list)"""

    expected_asserts_removed = """from code_to_optimize.bubble_sort import sorter
import pytest

def test_single_element_list():
    # Test that a single element list returns the same single element
    codeflash_output = sorter([42])

def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(1000, 0, -1))  # List from 1000 to 1
    codeflash_output = sorter(large_list)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    expected_instrumentation = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import pytest

def test_single_element_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sorter).bind([42])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_single_element_list', 'sorter', '1', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_single_element_list:sorter:1', codeflash_loop_index, codeflash_invocation_id)

def test_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_list = list(range(1000, 0, -1))
    _call__bound__arguments = inspect.signature(sorter).bind(large_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list', 'sorter', '2', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list:sorter:2', codeflash_loop_index, codeflash_invocation_id)"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    remove_asserts = remove_asserts_from_test(
        module=parse_module_to_cst(test_source),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    result = instrument_test_source(
        test_source=remove_asserts,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert remove_asserts == expected_asserts_removed
    assert result == expected_instrumentation


def test_instrument_class_method():
    test_source = """import pytest
from code_to_optimize.bubble_sort import Sorter

def test_single_element_list():
    # Test that a single element list returns the same single element
    obj = Sorter()
    codeflash_output = obj.sort([42])

def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(1000, 0, -1))  # List from 1000 to 1
    new_obj = Sorter()
    codeflash_output = new_obj.sort(large_list)
    """
    expected = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import Sorter
import pytest

def test_single_element_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    obj = Sorter()
    _call__bound__arguments = inspect.signature(obj.sort).bind([42])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(obj.sort, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_single_element_list', 'sort', '2', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_single_element_list:sort:2', codeflash_loop_index, codeflash_invocation_id)

def test_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_list = list(range(1000, 0, -1))
    new_obj = Sorter()
    _call__bound__arguments = inspect.signature(new_obj.sort).bind(large_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(new_obj.sort, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list', 'sort', '3', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list:sort:3', codeflash_loop_index, codeflash_invocation_id)"""
    function_to_optimize = FunctionToOptimize(
        function_name="sort",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[FunctionParent(name="Sorter", type="ClassDef")],
        starting_line=None,
        ending_line=None,
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert result == expected


def test_instrument_function_import_as():
    test_source = """import pytest
from code_to_optimize.bubble_sort import sorter as sort_me

def test_single_element_list():
    # Test that a single element list returns the same single element
    codeflash_output = sort_me([42])

def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(1000, 0, -1))  # List from 1000 to 1
    codeflash_output = sort_me(large_list)
    """

    expected = """import dill as pickle
import os
import gc
import time
def _log__test__values(values, duration, test_name, loop_index, invocation_id):
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with open(os.path.join("{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"), "ab") as f:
        if loop_index == 1:
            pickled_values = pickle.dumps(values)
        else:
            pickled_values = b""
        _test_name = f"{test_name}".encode("ascii")
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(duration.to_bytes(8, byteorder="big"))
        f.write(len(pickled_values).to_bytes(4, byteorder="big"))
        f.write(pickled_values)
        f.write(loop_index.to_bytes(8, byteorder="big"))
        f.write(len(invocation_id).to_bytes(4, byteorder="big"))
        f.write(invocation_id.encode("ascii"))

def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, *args, **kwargs):
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    print(f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}######!")
    gc.disable()
    counter = time.perf_counter_ns()
    return_value = wrapped(*args, **kwargs)
    codeflash_duration = time.perf_counter_ns() - counter
    gc.enable()
    return (return_value, codeflash_duration, invocation_id)
import inspect
from code_to_optimize.bubble_sort import sorter
import pytest
from code_to_optimize.bubble_sort import sorter as sort_me

def test_single_element_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sort_me).bind([42])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sort_me, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_single_element_list', 'sort_me', '1', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_single_element_list:sort_me:1', codeflash_loop_index, codeflash_invocation_id)

def test_large_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    large_list = list(range(1000, 0, -1))
    _call__bound__arguments = inspect.signature(sort_me).bind(large_list)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value, codeflash_duration, codeflash_invocation_id = codeflash_wrap(sort_me, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_list', 'sort_me', '2', codeflash_loop_index, **_call__bound__arguments.arguments)
    _log__test__values((_call__bound__arguments.arguments, codeflash_return_value), codeflash_duration, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0:test_large_list:sort_me:2', codeflash_loop_index, codeflash_invocation_id)"""
    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
    )

    assert result == expected
