from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
from core.languages.python.cst_utils import parse_module_to_cst
from core.languages.python.testgen.instrumentation.instrument_new_tests import instrument_test_source
from core.languages.python.testgen.instrumentation.wrapper import format_and_float_to_top
from core.languages.python.testgen.models import TestingMode
from core.languages.python.testgen.postprocessing.removeassert_transformer import remove_asserts_from_test
from tests.conftest import normalize_code

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

import dill as pickle


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f'{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}'
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_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_stdout_tag}######!')
    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'{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_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(codeflash_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


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f'{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}'
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_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_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value
"""


def test_instrument_async_test_source() -> None:
    test_source = """import pytest
import asyncio
from code_to_optimize.async_sorter import async_sorter

@pytest.mark.asyncio
async def test_single_element_list():
    # Test that a single element list returns the same single element
    codeflash_output = await async_sorter([42])

@pytest.mark.asyncio
async def test_large_list():
    # Test that a large list is sorted correctly
    large_list = list(range(10, 0, -1))  # List from 10 to 1

    codeflash_output = await async_sorter(large_list)

"""
    function_to_optimize = FunctionToOptimize(
        function_name="async_sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/async_sorter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
        is_async=True,
    )

    expected_result = """import asyncio
import inspect
import os

import pytest
from code_to_optimize.async_sorter import async_sorter


@pytest.mark.asyncio
async def test_single_element_list():
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    codeflash_output = await async_sorter([42])

@pytest.mark.asyncio
async def test_large_list():
    large_list = list(range(10, 0, -1))
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    codeflash_output = await async_sorter(large_list)"""

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.async_sorter",
        test_module_path="code_to_optimize.tests.pytest.test_async_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=30,
    )

    assert normalize_code(result) == normalize_code(expected_result)


def test_instrument_async_multiple_calls_same_test() -> None:
    """Test that multiple async function calls within the same test function get correctly numbered 0, 1, 2, etc."""
    test_source = """import pytest
import asyncio
from code_to_optimize.async_sorter import async_sorter

@pytest.mark.asyncio
async def test_single_call():
    result = await async_sorter([42])

@pytest.mark.asyncio
async def test_multiple_calls():
    result1 = await async_sorter([3, 1, 2])
    result2 = await async_sorter([5, 4])
    result3 = await async_sorter([9, 8, 7, 6])
    assert result1 == [1, 2, 3]
"""
    function_to_optimize = FunctionToOptimize(
        function_name="async_sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/async_sorter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
        is_async=True,
    )

    expected_result = """import asyncio
import inspect
import os

import pytest
from code_to_optimize.async_sorter import async_sorter


@pytest.mark.asyncio
async def test_single_call():
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    result = await async_sorter([42])

@pytest.mark.asyncio
async def test_multiple_calls():
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    result1 = await async_sorter([3, 1, 2])
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '1'
    result2 = await async_sorter([5, 4])
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '2'
    result3 = await async_sorter([9, 8, 7, 6])
    assert result1 == [1, 2, 3]"""

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.async_sorter",
        test_module_path="code_to_optimize.tests.pytest.test_async_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=30,
    )

    assert normalize_code(result) == normalize_code(expected_result)


def test_instrument_sync_vs_async_functions() -> None:
    """Test that sync and async functions get different instrumentation with exact string matching."""
    # Test SYNC function - should get full codeflash_wrap instrumentation
    sync_test_source = """import pytest
from code_to_optimize.bubble_sort import sorter

def test_sync_calls():
    result1 = sorter([3, 1, 2])
    result2 = sorter([5, 4])
"""

    sync_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,
        is_async=False,
    )

    expected_sync_result = """from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path

import dill as pickle
import pytest
from code_to_optimize.bubble_sort import sorter


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f'{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}'
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_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_stdout_tag}######!')
    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'{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_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(codeflash_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

def test_sync_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sorter).bind([3, 1, 2])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sorter, 'test_module', None, 'test_sync_calls', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result1 = codeflash_return_value
    _call__bound__arguments = inspect.signature(sorter).bind([5, 4])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sorter, 'test_module', None, 'test_sync_calls', 'sorter', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result2 = codeflash_return_value"""

    sync_result = instrument_test_source(
        test_source=sync_test_source,
        function_to_optimize=sync_function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
        test_module_path="test_module",
        test_framework="pytest",
        test_timeout=30,
    )

    assert normalize_code(sync_result) == normalize_code(expected_sync_result)

    # Test ASYNC function - should get os.environ assignments only
    async_test_source = """import pytest
from code_to_optimize.async_sorter import async_sorter

@pytest.mark.asyncio
async def test_async_calls():
    result1 = await async_sorter([3, 1, 2])
    result2 = await async_sorter([5, 4])"""

    async_function_to_optimize = FunctionToOptimize(
        function_name="async_sorter",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/async_sorter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
        is_async=True,
    )

    expected_async_result = """import inspect
import os

import pytest
from code_to_optimize.async_sorter import async_sorter


@pytest.mark.asyncio
async def test_async_calls():
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    result1 = await async_sorter([3, 1, 2])
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '1'
    result2 = await async_sorter([5, 4])"""

    async_result = instrument_test_source(
        test_source=async_test_source,
        function_to_optimize=async_function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.async_sorter",
        test_module_path="test_module",
        test_framework="pytest",
        test_timeout=30,
    )

    assert async_result == expected_async_result


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(
        "\n"
        + 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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value"""
    )
    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 normalize_code(result) == normalize_code(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])
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value"""
    )
    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 normalize_code(result) == normalize_code(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(
        behavior_logging_code
        + """import pytest
from code_to_optimize.bubble_sort_in_class import BubbleSortClass
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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value"""
    )
    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 normalize_code(result) == normalize_code(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])
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value"""
    )
    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 normalize_code(result) == normalize_code(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 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')
    codeflash_output = codeflash_return_value

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')
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value

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])
    codeflash_output = codeflash_return_value"""
    )
    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 normalize_code(result) == normalize_code(expected)


def test_instrument_test_source4() -> None:
    test_source = """import pytest
from code_to_optimize.bubble_sort import sorter
class Config:
    xy_point_dist = 1.0
    xy_thickness = 0.1

@pytest.fixture
def processor():
    return GCodeProcessor(Config())


def test_minimal_distance(processor):
    # Test with very close start and end points
    start = (0, 0, 0, 0)
    end = (0.001, 0.001, 0, 0.001)
    result = processor.interpolate_with_constant_resolution_XY(start, end, 0.001)
    assert len(result) == 2  # Should still have at least start and end points

def test_wobble_application(processor):
    # Test wobble is applied correctly
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    result = processor.interpolate_with_constant_resolution_XY(start, end, 10)
    # Ensure no wobble at start and end
    assert result[0] == (0, 0, 0, 0)
    assert result[-1] == (10, 0, 0, 1)
    # Check that wobble is applied to intermediate points
    assert result[1][0] != 1  # x should be wobbled

def test_extrusion_handling(processor):
    # Test extrusion when start point includes extrusion
    start = (0, 0, 0, 5)
    end = (10, 0, 0, 15)
    result = processor.interpolate_with_constant_resolution_XY(start, end, 10)
    assert result[0] == (0, 0, 0, 5)  # Start point with absolute extrusion
    assert result[-1] == (10, 0, 0, 1)  # End point with delta extrusion

def test_large_distance_interpolation(processor):
    # Test long distance interpolation
    start = (0, 0, 0, 0)
    end = (1000, 0, 0, 1000)
    result = processor.interpolate_with_constant_resolution_XY(start, end, 1000)
    assert len(result) == 1001  # 1000 segments + 1 start point

def test_high_number_of_points(processor):
    # Test with high number of points due to small xy_point_dist
    processor.config.xy_point_dist = 0.01
    start = (0, 0, 0, 0)
    end = (1, 0, 0, 1)
    result = processor.interpolate_with_constant_resolution_XY(start, end, 1)
    assert len(result) == 101  # 100 segments + 1 start point

def test_consistent_results_with_fixed_seed(processor):
    # Test consistent results with a fixed random seed
    random.seed(42)
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    result1 = processor.interpolate_with_constant_resolution_XY(start, end, 10)
    result2 = processor.interpolate_with_constant_resolution_XY(start, end, 10)
    assert result1 == result2  # Results should be identical

def test_consistent_results_with_tuple_ret(processor):
    # Test consistent results with a fixed random seed
    random.seed(42)
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    result1, result2 = processor.interpolate_with_constant_resolution_XY(start, end, 10)
    assert result1 == result2  # Results should be identical"""

    function_to_optimize = FunctionToOptimize(
        function_name="interpolate_with_constant_resolution_XY",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    test_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 = format_and_float_to_top(
        behavior_logging_code
        + """import pytest
from code_to_optimize.bubble_sort import interpolate_with_constant_resolution_XY, sorter

class Config:
    xy_point_dist = 1.0
    xy_thickness = 0.1

@pytest.fixture
def processor():
    return GCodeProcessor(Config())

def test_minimal_distance(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    start = (0, 0, 0, 0)
    end = (0.001, 0.001, 0, 0.001)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 0.001)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_minimal_distance', 'interpolate_with_constant_resolution_XY', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert len(result) == 2

def test_wobble_application(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_wobble_application', 'interpolate_with_constant_resolution_XY', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result[0] == (0, 0, 0, 0)
    assert result[-1] == (10, 0, 0, 1)
    assert result[1][0] != 1

def test_extrusion_handling(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    start = (0, 0, 0, 5)
    end = (10, 0, 0, 15)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_extrusion_handling', 'interpolate_with_constant_resolution_XY', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result[0] == (0, 0, 0, 5)
    assert result[-1] == (10, 0, 0, 1)

def test_large_distance_interpolation(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    start = (0, 0, 0, 0)
    end = (1000, 0, 0, 1000)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 1000)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_large_distance_interpolation', 'interpolate_with_constant_resolution_XY', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert len(result) == 1001

def test_high_number_of_points(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    processor.config.xy_point_dist = 0.01
    start = (0, 0, 0, 0)
    end = (1, 0, 0, 1)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 1)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_high_number_of_points', 'interpolate_with_constant_resolution_XY', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert len(result) == 101

def test_consistent_results_with_fixed_seed(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    random.seed(42)
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_consistent_results_with_fixed_seed', 'interpolate_with_constant_resolution_XY', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result1 = codeflash_return_value
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_consistent_results_with_fixed_seed', 'interpolate_with_constant_resolution_XY', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result2 = codeflash_return_value
    assert result1 == result2

def test_consistent_results_with_tuple_ret(processor):
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    random.seed(42)
    start = (0, 0, 0, 0)
    end = (10, 0, 0, 10)
    _call__bound__arguments = inspect.signature(processor.interpolate_with_constant_resolution_XY).bind(start, end, 10)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(processor.interpolate_with_constant_resolution_XY, 'code_to_optimize.tests.pytest.test_sorter__unit_test_0', None, 'test_consistent_results_with_tuple_ret', 'interpolate_with_constant_resolution_XY', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result1, result2 = codeflash_return_value
    assert result1 == result2"""
    )

    assert normalize_code(test_result) == normalize_code(expected_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(
#         behavior_logging_code
#         + """
# import pytest
# from code_to_optimize.bubble_sort_in_class import BubbleSortClass
# 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.args, **_call__bound__arguments.kwargs)
#     codeflash_output = codeflash_return_value

# 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.args, **_call__bound__arguments.kwargs)
#     codeflash_output = codeflash_return_value"""
#     )
#     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 normalize_code(result) == normalize_code(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])
#     codeflash_output = codeflash_return_value

# 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)
#     codeflash_output = codeflash_return_value"""
#     )
#     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 normalize_code(result) == normalize_code(expected_perfonly_result)


def test_instrument_random_test() -> None:
    #     test_source = """import numpy as np
    # def test_np_np_sum():
    #     return np.sum(a + b)
    # """

    #     expected_result = format_and_float_to_top(
    #         behavior_logging_code
    #         + """import numpy as np
    # from code_to_optimize.bubble_sort import np_np_sum
    # def test_np_np_sum():
    #     np.random.seed(42)
    #     codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    #     return np.sum(a + b)"""
    #     )

    #     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 normalize_code(result) == normalize_code(expected_result)
    function_to_optimize = FunctionToOptimize(
        function_name="np_np_sum",
        file_path="/Users/renaud/repos/codeflash/cli/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    test_source2 = """import numpy
def test_np_np_sum():
    return numpy.sum(a + b)
"""
    expected_result2 = format_and_float_to_top(
        behavior_logging_code
        + """
import numpy
from code_to_optimize.bubble_sort import np_np_sum
def test_np_np_sum():
    numpy.random.seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    return numpy.sum(a + b)
"""
    )

    result2 = instrument_test_source(
        test_source=test_source2,
        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 normalize_code(result2) == normalize_code(expected_result2)

    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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )
    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 normalize_code(result) == normalize_code(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)
    codeflash_output = codeflash_return_value
"""
    )
    assert normalize_code(result) == normalize_code(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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )
    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 normalize_code(result) == normalize_code(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)
    codeflash_output = codeflash_return_value
"""
    )
    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 normalize_code(result) == normalize_code(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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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 normalize_code(result) == normalize_code(
        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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )

    assert normalize_code(result) == normalize_code(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)
    codeflash_output = codeflash_return_value
"""
    )

    assert normalize_code(result) == normalize_code(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.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
"""
    )

    assert normalize_code(result) == normalize_code(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)
        codeflash_output = codeflash_return_value
"""
    )

    assert normalize_code(result) == normalize_code(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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )

    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 normalize_code(result) == normalize_code(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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )

    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 normalize_code(result) == normalize_code(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])
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value
"""
    )

    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 normalize_code(result) == normalize_code(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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

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.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
"""
    )

    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 normalize_code(result) == normalize_code(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])
    codeflash_output = codeflash_return_value

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)
    codeflash_output = codeflash_return_value
"""
    )

    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 normalize_code(result) == normalize_code(expected)


def test_instrument_list_comprehension_skips_instrumentation() -> None:
    """Test that function calls inside list comprehensions are NOT instrumented.

    This prevents NameError when the loop variable (e.g., 'ts') is only valid
    inside the comprehension scope but would be used outside after extraction.
    """
    test_source = """from code_to_optimize.converter import convert_item

def test_list_comprehension():
    items = ['a', 'b', 'c']
    codeflash_output = [convert_item(item) for item in items]
"""

    function_to_optimize = FunctionToOptimize(
        function_name="convert_item",
        file_path="/Users/test/code_to_optimize/converter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    # The list comprehension should NOT be instrumented - the call should remain inside
    expected = format_and_float_to_top(
        perf_logging_code
        + """
from code_to_optimize.converter import convert_item
def test_list_comprehension():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    items = ['a', 'b', 'c']
    codeflash_output = [convert_item(item) for item in items]
"""
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.converter",
        test_module_path="code_to_optimize.tests.pytest.test_converter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.PERFORMANCE,
    )

    # Normalize via ast.unparse to ignore quote style differences
    assert normalize_code(result) == normalize_code(expected)


def test_instrument_generator_expression_skips_instrumentation() -> None:
    """Test that function calls inside generator expressions are NOT instrumented."""
    test_source = """from code_to_optimize.converter import convert_item

def test_generator_expression():
    items = ['a', 'b', 'c']
    codeflash_output = list(convert_item(item) for item in items)
"""

    function_to_optimize = FunctionToOptimize(
        function_name="convert_item",
        file_path="/Users/test/code_to_optimize/converter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    # The generator expression should NOT be instrumented
    expected = format_and_float_to_top(
        perf_logging_code
        + """
from code_to_optimize.converter import convert_item
def test_generator_expression():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    items = ['a', 'b', 'c']
    codeflash_output = list((convert_item(item) for item in items))
"""
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.converter",
        test_module_path="code_to_optimize.tests.pytest.test_converter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.PERFORMANCE,
    )

    # Normalize via ast.unparse to ignore quote style differences
    assert normalize_code(result) == normalize_code(expected)


def test_instrument_dict_comprehension_skips_instrumentation() -> None:
    """Test that function calls inside dict comprehensions are NOT instrumented."""
    test_source = """from code_to_optimize.converter import convert_item

def test_dict_comprehension():
    items = ['a', 'b', 'c']
    codeflash_output = {item: convert_item(item) for item in items}
"""

    function_to_optimize = FunctionToOptimize(
        function_name="convert_item",
        file_path="/Users/test/code_to_optimize/converter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    # The dict comprehension should NOT be instrumented
    expected = format_and_float_to_top(
        perf_logging_code
        + """
from code_to_optimize.converter import convert_item
def test_dict_comprehension():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    items = ['a', 'b', 'c']
    codeflash_output = {item: convert_item(item) for item in items}
"""
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.converter",
        test_module_path="code_to_optimize.tests.pytest.test_converter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.PERFORMANCE,
    )

    # Normalize via ast.unparse to ignore quote style differences
    assert normalize_code(result) == normalize_code(expected)


def test_instrument_removes_import_inspect_inside_function() -> None:
    """Test that 'import inspect' inside a test function is removed to avoid UnboundLocalError.

    When the LLM generates tests with `import inspect` inside a function body (e.g., before
    an assert that gets removed by postprocessing), we need to remove that import because:
    1. Python treats `import X` inside a function as making X a local variable
    2. The instrumentation code uses `inspect.signature()` BEFORE the import would execute
    3. This causes UnboundLocalError: cannot access local variable 'inspect'

    The fix is to skip `import inspect` inside function bodies since it's added at module level.
    """
    # This simulates a test where the LLM generated `import inspect` inside the function
    # (perhaps before an assert like `assert inspect.iscoroutinefunction(wrapper)` that was removed)
    test_source = """from code_to_optimize.bubble_sort import sorter

def test_with_import_inspect_inside():
    result = sorter([3, 1, 2])
    import inspect  # This would cause UnboundLocalError if not removed
"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="/Users/test/code_to_optimize/bubble_sort.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    # Expected: import inspect is at module level, NOT inside the function
    # The function body should have instrumentation code but no `import inspect`
    expected = format_and_float_to_top(
        behavior_logging_code
        + """
from code_to_optimize.bubble_sort import sorter
def test_with_import_inspect_inside():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(sorter).bind([3, 1, 2])
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize.tests.pytest.test_bubble_sort__unit_test_0', None, 'test_with_import_inspect_inside', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
"""
    )

    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_bubble_sort__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.BEHAVIOR,
    )

    assert normalize_code(result) == normalize_code(expected)


def test_instrument_removes_import_os_inside_async_function() -> None:
    """Test that 'import os' inside an async test function is removed to avoid UnboundLocalError.

    For async functions, instrumentation adds `os.environ['CODEFLASH_CURRENT_LINE_ID'] = ...`
    before the function call. If there's also `import os` inside the function body,
    Python treats `os` as a local variable and causes UnboundLocalError.
    """
    test_source = """import pytest
from code_to_optimize.async_sorter import async_sorter

@pytest.mark.asyncio
async def test_with_import_os_inside():
    result = await async_sorter([3, 1, 2])
    import os  # This would cause UnboundLocalError if not removed
"""

    function_to_optimize = FunctionToOptimize(
        function_name="async_sorter",
        file_path="/Users/test/code_to_optimize/async_sorter.py",
        parents=[],
        starting_line=None,
        ending_line=None,
        is_async=True,
    )

    # Expected: import os is at module level, NOT inside the function
    # The function body should have os.environ set but no `import os`
    expected = """import inspect
import os

import pytest
from code_to_optimize.async_sorter import async_sorter


@pytest.mark.asyncio
async def test_with_import_os_inside():
    os.environ['CODEFLASH_CURRENT_LINE_ID'] = '0'
    result = await async_sorter([3, 1, 2])"""

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.async_sorter",
        test_module_path="code_to_optimize.tests.pytest.test_async_sorter__unit_test_0",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.BEHAVIOR,
    )

    assert normalize_code(result) == normalize_code(expected)


def test_instrument_chained_call_decorator_pattern() -> None:
    """Test that chained calls like action_wrap(X)(Y) are instrumented correctly.

    When a function returns a callable that is immediately called (decorator pattern),
    the instrumentation should capture the result of the outermost call, not the
    intermediate decorator function.

    For example, with `action_wrap(action)(fn)`:
    - action_wrap(action) returns a decorator
    - decorator(fn) returns the wrapped function
    We want to capture the wrapped function, not the decorator.
    """
    test_source = """import pytest
from code_to_optimize.decorators import action_wrap

def test_chained_call():
    wrapper = action_wrap("click")(lambda x: x)
    result = wrapper("test")
"""

    function_to_optimize = FunctionToOptimize(
        function_name="action_wrap",
        file_path="/Users/test/code_to_optimize/decorators.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )

    # The key assertion: the instrumentation should wrap the OUTER call (action_wrap("click")(lambda x: x))
    # not just the inner call (action_wrap("click"))
    expected_result = format_and_float_to_top(
        behavior_logging_code
        + """import pytest
from code_to_optimize.decorators import action_wrap

def test_chained_call():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(action_wrap('click')).bind(lambda x: x)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(action_wrap('click'), 'code_to_optimize.tests.pytest.test_action_wrap', None, 'test_chained_call', 'action_wrap', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    wrapper = codeflash_return_value
    result = wrapper('test')"""
    )

    result = instrument_test_source(
        test_source=test_source,
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.decorators",
        test_module_path="code_to_optimize.tests.pytest.test_action_wrap",
        test_framework="pytest",
        test_timeout=15,
        mode=TestingMode.BEHAVIOR,
    )

    assert normalize_code(result) == normalize_code(expected_result)
