codeflash-internal/django/aiservice/testgen/instrumentation/tests/test_instrument_generated_tests.py

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 format_and_float_to_top, instrument_test_source
from testgen.models import TestingMode
from testgen.postprocessing.removeassert_transformer import remove_asserts_from_test

imports_block = """from __future__ import annotations

import gc
import os
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable, Optional


"""
behavior_logging_code = """from __future__ import annotations

import gc
import inspect
import os
import time

from pathlib import Path
from typing import Any, Callable, Optional

import dill as pickle


def codeflash_wrap(
    wrapped: Callable[..., Any],
    test_module_name: str,
    test_class_name: str | None,
    test_name: str,
    function_name: str,
    line_id: str,
    loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    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}######!"
    )
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    iteration = os.environ["CODEFLASH_TEST_ITERATION"]
    with Path(
        "{codeflash_run_tmp_dir_client_side}", f"test_return_values_{iteration}.bin"
    ).open("ab") as f:
        pickled_values = (
            pickle.dumps((args, kwargs, exception))
            if exception
            else pickle.dumps((args, kwargs, return_value))
        )
        _test_name = f"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{line_id}".encode(
            "ascii"
        )
        f.write(len(_test_name).to_bytes(4, byteorder="big"))
        f.write(_test_name)
        f.write(codeflash_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"))
    if exception:
        raise exception
    return return_value
"""

perf_logging_code = """from __future__ import annotations

import gc
import os
import time
from typing import Any, Callable, Optional


def codeflash_wrap(
    wrapped: Callable[..., Any],
    test_module_name: str,
    test_class_name: Optional[str],
    test_name: str,
    function_name: str,
    line_id: str,
    loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(
        f"!######{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}:{codeflash_duration}######!"
    )
    if exception:
        raise exception
    return return_value
"""


def test_instrument_test_source() -> None:
    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 = format_and_float_to_top(
        behavior_logging_code  # TODO : There should not be 4 blank lines here, fix this in the instrumentation
        + """import pytest
from code_to_optimize.bubble_sort import sorter


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sorter",
        "0",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sorter",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )
    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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == expected_result

    expected_perfonly_result = format_and_float_to_top(
        perf_logging_code
        + """

from code_to_optimize.bubble_sort import sorter
import pytest


def test_single_element_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    codeflash_return_value = codeflash_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sorter",
        "0",
        codeflash_loop_index,
        [42],
    )


def test_large_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    large_list = list(range(1000, 0, -1))
    codeflash_return_value = codeflash_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sorter",
        "1",
        codeflash_loop_index,
        large_list,
    )
"""
    )
    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,
        mode=TestingMode.PERFORMANCE,
    )
    assert result == expected_perfonly_result


def test_instrument_test_source2() -> None:
    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 = format_and_float_to_top(
        """import pytest
from code_to_optimize.bubble_sort_in_class import BubbleSortClass
"""
        + behavior_logging_code
        + """


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_wrap(
        bubble_sort_instance.sorter,
        "tests.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sorter",
        "0",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


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_wrap(
        bubble_sort_instance.sorter,
        "tests.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sorter",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )
    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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == expected_result

    expected_perfonly_result = format_and_float_to_top(
        perf_logging_code
        + """

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"])
    codeflash_return_value = codeflash_wrap(
        bubble_sort_instance.sorter,
        "tests.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sorter",
        "0",
        codeflash_loop_index,
        [42],
    )


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))
    codeflash_return_value = codeflash_wrap(
        bubble_sort_instance.sorter,
        "tests.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sorter",
        "2",
        codeflash_loop_index,
        large_list,
    )
"""
    )
    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,
        mode=TestingMode.PERFORMANCE,
    )
    assert result == expected_perfonly_result


def test_instrument_test_source3() -> None:
    test_source = """from __future__ import annotations

import typing as t
from unittest.mock import Mock

# imports
import pytest  # used for our unit tests
from src.click.core import CommandCollection


class Command:
    pass  # Placeholder for the Command class

class Context:
    pass  # Placeholder for the Context class
from src.click.core import CommandCollection

# unit tests

@pytest.fixture
def mock_context():
    return Mock(spec=Context)

@pytest.fixture
def mock_command():
    return Mock(spec=Command)

@pytest.fixture
def mock_group():
    group = Mock(spec=Group)
    group.get_command.return_value = None
    return group

def test_command_exists_in_current_collection(mock_context, mock_command):
    # Setup
    collection = CommandCollection()
    collection.get_command = Mock(return_value=mock_command)

    # Test
    codeflash_output = collection.get_command(mock_context, "start")

def test_command_exists_in_source_group(mock_context, mock_command, mock_group):
    # Setup
    mock_group.get_command.return_value = mock_command
    collection = CommandCollection(sources=[mock_group])

    # Test
    codeflash_output = collection.get_command(mock_context, option="stop")

def test_command_does_not_exist(mock_context, mock_group):
    # Setup
    collection = CommandCollection(sources=[mock_group])
    args = ["start", "stop", "restart"]
    kwargs = {"option2": "stop"}

    # Test
    codeflash_output = collection.get_command(mock_context, *args, option="restart", **kwargs)

def test_empty_command_name(mock_context, mock_group):
    # Setup
    collection = CommandCollection(sources=[mock_group])

    # Test
    codeflash_output = collection.get_command(mock_context)

def test_none_as_command_name(mock_context, mock_group):
    # Setup
    collection = CommandCollection(sources=[mock_group])

    # Test
    codeflash_output = collection.get_command(mock_context, sources=[mock_group])
"""

    expected = format_and_float_to_top(
        """from __future__ import annotations

import gc
import os
import time
import typing as t
from typing import Any, Callable, Optional
from unittest.mock import Mock

import pytest
from code_to_optimize.bubble_sort_in_class import CommandCollection

"""
        + perf_logging_code
        + """


class Command:
    pass


class Context:
    pass


@pytest.fixture
def mock_context():
    return Mock(spec=Context)


@pytest.fixture
def mock_command():
    return Mock(spec=Command)


@pytest.fixture
def mock_group():
    group = Mock(spec=Group)
    group.get_command.return_value = None
    return group


def test_command_exists_in_current_collection(mock_context, mock_command):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    collection = CommandCollection()
    collection.get_command = Mock(return_value=mock_command)
    codeflash_return_value = codeflash_wrap(
        collection.get_command,
        "tests.test_sorter__unit_test_0",
        None,
        "test_command_exists_in_current_collection",
        "get_command",
        "2",
        codeflash_loop_index,
        mock_context,
        "start",
    )


def test_command_exists_in_source_group(mock_context, mock_command, mock_group):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    mock_group.get_command.return_value = mock_command
    collection = CommandCollection(sources=[mock_group])
    codeflash_return_value = codeflash_wrap(
        collection.get_command,
        "tests.test_sorter__unit_test_0",
        None,
        "test_command_exists_in_source_group",
        "get_command",
        "2",
        codeflash_loop_index,
        mock_context,
        option="stop",
    )


def test_command_does_not_exist(mock_context, mock_group):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    collection = CommandCollection(sources=[mock_group])
    args = ["start", "stop", "restart"]
    kwargs = {"option2": "stop"}
    codeflash_return_value = codeflash_wrap(
        collection.get_command,
        "tests.test_sorter__unit_test_0",
        None,
        "test_command_does_not_exist",
        "get_command",
        "3",
        codeflash_loop_index,
        mock_context,
        *args,
        option="restart",
        **kwargs,
    )


def test_empty_command_name(mock_context, mock_group):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    collection = CommandCollection(sources=[mock_group])
    codeflash_return_value = codeflash_wrap(
        collection.get_command,
        "tests.test_sorter__unit_test_0",
        None,
        "test_empty_command_name",
        "get_command",
        "1",
        codeflash_loop_index,
        mock_context,
    )


def test_none_as_command_name(mock_context, mock_group):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    collection = CommandCollection(sources=[mock_group])
    codeflash_return_value = codeflash_wrap(
        collection.get_command,
        "tests.test_sorter__unit_test_0",
        None,
        "test_none_as_command_name",
        "get_command",
        "1",
        codeflash_loop_index,
        mock_context,
        sources=[mock_group],
    )
"""
    )
    function_to_optimize = FunctionToOptimize(
        function_name="get_command",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort_in_class.py",
        parents=[FunctionParent(name="CommandCollection", 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_in_class",
        test_module_path="tests.test_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.PERFORMANCE,
    )
    assert result == expected


def test_instrument_random_test() -> None:
    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 = format_and_float_to_top(
        behavior_logging_code
        + """
import random
from code_to_optimize.bubble_sort import sorter


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list_with_random_values",
        "sorter",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )
    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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == expected_result

    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,
        mode=TestingMode.PERFORMANCE,
    )
    expected_perf_result = format_and_float_to_top(
        perf_logging_code
        + """

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)
    codeflash_return_value = codeflash_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list_with_random_values",
        "sorter",
        "2",
        codeflash_loop_index,
        random_list,
    )
"""
    )
    assert result == expected_perf_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 = format_and_float_to_top(
        behavior_logging_code
        + """
import random
from random import sample
from code_to_optimize.bubble_sort import sorter


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list_with_random_values",
        "sorter",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )
    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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == expected_result

    expected_result = format_and_float_to_top(
        perf_logging_code
        + """

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)
    codeflash_return_value = codeflash_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list_with_random_values",
        "sorter",
        "2",
        codeflash_loop_index,
        random_list,
    )
"""
    )
    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,
        mode=TestingMode.PERFORMANCE,
    )

    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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """

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 = """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,
        mode=TestingMode.PERFORMANCE,
    )

    assert result == format_and_float_to_top(
        perf_logging_code
        + """

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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """

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
"""
    )

    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,
        mode=TestingMode.PERFORMANCE,
    )

    assert result == format_and_float_to_top(
        perf_logging_code
        + """

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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
import random
from code_to_optimize.bubble_sort import sorter


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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
import numpy.random
from code_to_optimize.bubble_sort import sorter


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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )
    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
import random
from code_to_optimize.bubble_sort import sorter


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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == format_and_float_to_top(
        behavior_logging_code
        + """
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() -> None:
    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,
        mode=TestingMode.BEHAVIOR,
    )
    expected_result = format_and_float_to_top(
        behavior_logging_code
        + """
import pytest
from code_to_optimize.bubble_sort import sorter


@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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_sort_parametrized",
        "sorter",
        "0",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )

    assert result == expected_result

    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,
        mode=TestingMode.PERFORMANCE,
    )
    expected_result = format_and_float_to_top(
        perf_logging_code
        + """

from code_to_optimize.bubble_sort import sorter
import pytest


@pytest.mark.parametrize(
    "input, expected_output",
    [
        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
        (list(reversed(range(50))), list(range(50))),
    ],
)
def test_sort_parametrized(input, expected_output):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    codeflash_return_value = codeflash_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_sort_parametrized",
        "sorter",
        "0",
        codeflash_loop_index,
        input,
    )
"""
    )

    assert result == expected_result


def test_instrument_parameterized_test_with_loop() -> None:
    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,
        mode=TestingMode.BEHAVIOR,
    )

    expected_result = format_and_float_to_top(
        behavior_logging_code
        + """
from code_to_optimize.bubble_sort import sorter
import pytest


@pytest.mark.parametrize(
    "input, expected_output",
    [
        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
        (list(reversed(range(50))), list(range(50))),
    ],
)
def test_sort_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_wrap(
            sorter,
            "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
            None,
            "test_sort_loop_parametrized",
            "sorter",
            "0_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
"""
    )

    assert result == expected_result

    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,
        mode=TestingMode.PERFORMANCE,
    )

    expected_result = format_and_float_to_top(
        perf_logging_code
        + """

from code_to_optimize.bubble_sort import sorter
import pytest


@pytest.mark.parametrize(
    "input, expected_output",
    [
        ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]),
        ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]),
        (list(reversed(range(50))), list(range(50))),
    ],
)
def test_sort_loop_parametrized(input, expected_output):
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    for i in range(2):
        codeflash_return_value = codeflash_wrap(
            sorter,
            "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
            None,
            "test_sort_loop_parametrized",
            "sorter",
            "0_0",
            codeflash_loop_index,
            input,
        )
"""
    )

    assert result == expected_result


def test_conditional_instrumentation() -> None:
    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,
        mode=TestingMode.BEHAVIOR,
    )

    expected_result = format_and_float_to_top(
        behavior_logging_code
        + """
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_wrap(
            sorter,
            "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
            None,
            "test_sort",
            "sorter",
            "1_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )
"""
    )

    assert result == expected_result

    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,
        mode=TestingMode.PERFORMANCE,
    )

    expected_result = format_and_float_to_top(
        perf_logging_code
        + """

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:
        codeflash_return_value = codeflash_wrap(
            sorter,
            "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
            None,
            "test_sort",
            "sorter",
            "1_0",
            codeflash_loop_index,
            input,
        )
"""
    )

    assert result == expected_result


def test_remove_asserts_and_instrument_test_source() -> None:
    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 = format_and_float_to_top(
        behavior_logging_code
        + """
import pytest
from code_to_optimize.bubble_sort import sorter


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sorter",
        "0",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


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_wrap(
        sorter,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sorter",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )

    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,
        mode=TestingMode.BEHAVIOR,
    )

    assert remove_asserts == expected_asserts_removed
    assert result == expected_instrumentation


def test_instrument_class_method() -> None:
    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)
    """

    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,
    )
    expected = format_and_float_to_top(
        behavior_logging_code
        + """
import pytest
from code_to_optimize.bubble_sort import Sorter


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_wrap(
        obj.sort,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sort",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


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_wrap(
        new_obj.sort,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sort",
        "2",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )

    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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == expected

    expected = format_and_float_to_top(
        perf_logging_code
        + """

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()
    codeflash_return_value = codeflash_wrap(
        obj.sort,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sort",
        "1",
        codeflash_loop_index,
        [42],
    )


def test_large_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    large_list = list(range(1000, 0, -1))
    new_obj = Sorter()
    codeflash_return_value = codeflash_wrap(
        new_obj.sort,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sort",
        "2",
        codeflash_loop_index,
        large_list,
    )
"""
    )

    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,
        mode=TestingMode.PERFORMANCE,
    )

    assert result == expected


def test_instrument_function_import_as() -> None:
    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)
"""
    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 = format_and_float_to_top(
        behavior_logging_code
        + """
import pytest
from code_to_optimize.bubble_sort import sorter
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_wrap(
        sort_me,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sort_me",
        "0",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )


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_wrap(
        sort_me,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sort_me",
        "1",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
"""
    )

    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,
        mode=TestingMode.BEHAVIOR,
    )

    assert result == expected

    expected = format_and_float_to_top(
        perf_logging_code
        + """

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"])
    codeflash_return_value = codeflash_wrap(
        sort_me,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_single_element_list",
        "sort_me",
        "0",
        codeflash_loop_index,
        [42],
    )


def test_large_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    large_list = list(range(1000, 0, -1))
    codeflash_return_value = codeflash_wrap(
        sort_me,
        "code_to_optimize.tests.pytest.test_sorter__unit_test_0",
        None,
        "test_large_list",
        "sort_me",
        "1",
        codeflash_loop_index,
        large_list,
    )
"""
    )

    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,
        mode=TestingMode.PERFORMANCE,
    )

    assert result == expected