import ast
import os

from core.languages.python.cst_utils import parse_module_to_cst
from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
from core.languages.python.testgen.instrumentation.instrument_new_tests import InjectPerfAndLogging
from core.languages.python.testgen.models import TestingMode
from core.languages.python.testgen.postprocessing.removeassert_transformer import remove_asserts_from_test

os.environ["CODEFLASH_API_KEY"] = "cf-test-key"


def test_InjectPerfAndLogging_with() -> None:
    code = """def test_relative_validity_no_tree():
    hdbscan = HDBSCAN()
    with pytest.raises(AttributeError):
        hdbscan.relative_validity_()"""
    auxillary_functions: list[str] = []
    module_node = ast.parse(code)
    test_module_path = "code_to_optimize_path"
    function_to_optimize = FunctionToOptimize(
        function_name="relative_validity_",
        file_path="/tmp/path",
        parents=[FunctionParent(name="HDBSCAN", type="ClassDef")],
    )
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.BEHAVIOR
    ).visit(module_node)

    expected = """def test_relative_validity_no_tree():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    hdbscan = HDBSCAN()
    with pytest.raises(AttributeError):
        _call__bound__arguments = inspect.signature(hdbscan.relative_validity_).bind()
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)"""
    assert ast.unparse(new_module_node).strip("\n") == expected
    module_node = ast.parse(code)
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.PERFORMANCE
    ).visit(module_node)

    expected = """def test_relative_validity_no_tree():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    hdbscan = HDBSCAN()
    with pytest.raises(AttributeError):
        codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1_0', codeflash_loop_index)"""
    assert ast.unparse(new_module_node).strip("\n") == expected


def test_InjectPerfAndLogging() -> None:
    code = """def test_relative_validity_no_tree():
    hdbscan = HDBSCAN()
    hdbscan.relative_validity_()"""
    auxillary_functions: list[str] = []
    module_node = ast.parse(code)
    test_module_path = "code_to_optimize_path"
    function_to_optimize = FunctionToOptimize(
        function_name="relative_validity_",
        file_path="/tmp/path",
        parents=[FunctionParent(name="HDBSCAN", type="ClassDef")],
    )
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.BEHAVIOR
    ).visit(module_node)
    expected = """def test_relative_validity_no_tree():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    hdbscan = HDBSCAN()
    _call__bound__arguments = inspect.signature(hdbscan.relative_validity_).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)"""
    assert ast.unparse(new_module_node).strip("\n") == expected
    module_node = ast.parse(code)
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.PERFORMANCE
    ).visit(module_node)
    expected = """def test_relative_validity_no_tree():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    hdbscan = HDBSCAN()
    codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1', codeflash_loop_index)"""
    assert ast.unparse(new_module_node).strip("\n") == expected


def test_remove_bad_assert() -> None:
    code = """def test_relative_validity_no_tree():
    hdbscan = HDBSCAN()
    codeflash_output = hdbscan.relative_validity_() == 0.5
    result = 5"""
    auxillary_functions: list[str] = []
    module_node = ast.parse(code)
    test_module_path = "code_to_optimize_path"
    function_to_optimize = FunctionToOptimize(
        function_name="relative_validity_",
        file_path="/tmp/path",
        parents=[FunctionParent(name="HDBSCAN", type="ClassDef")],
    )
    new_module_node = InjectPerfAndLogging(function_to_optimize, auxillary_functions, test_module_path).visit(
        module_node
    )
    expected = """def test_relative_validity_no_tree():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    hdbscan = HDBSCAN()
    _call__bound__arguments = inspect.signature(hdbscan.relative_validity_).bind()
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(hdbscan.relative_validity_, 'code_to_optimize_path', None, 'test_relative_validity_no_tree', 'relative_validity_', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = 5"""
    assert ast.unparse(new_module_node).strip("\n") == expected

    code = """def test_translate_word_starting_with_vowel():
    codeflash_output = translate('apple')
def test_translate_word_starting_with_single_consonant():
    codeflash_output = translate('banana')"""

    function_to_optimize = FunctionToOptimize(function_name="translate", file_path="/tmp/path", parents=[])
    auxillary_functions = []
    module_node = ast.parse(code)
    test_module_path = "code_to_optimize_path"
    new_module_node = InjectPerfAndLogging(function_to_optimize, auxillary_functions, test_module_path).visit(
        module_node
    )
    expected = """def test_translate_word_starting_with_vowel():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(translate).bind('apple')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(translate, 'code_to_optimize_path', None, 'test_translate_word_starting_with_vowel', 'translate', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value

def test_translate_word_starting_with_single_consonant():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(translate).bind('banana')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(translate, 'code_to_optimize_path', None, 'test_translate_word_starting_with_single_consonant', 'translate', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value"""
    assert ast.unparse(new_module_node).strip("\n") == expected


def test_unittest_generated_tests_bubble_sort() -> None:
    original_test = """import unittest
def sorter(arr):
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp
    return arr
class SorterTestCase(unittest.TestCase):
    def test_empty_list(self):
        self.assertEqual(sorter([]), [])
    def test_single_element_list(self):
        self.assertEqual(sorter([5]), [5])
    def test_ascending_order_list(self):
        self.assertEqual(sorter([1, 2, 3, 4, 5]), [1, 2, 3, 4, 5])
    def test_descending_order_list(self):
        self.assertEqual(sorter([5, 4, 3, 2, 1]), [1, 2, 3, 4, 5])
    def test_random_order_list(self):
        self.assertEqual(sorter([3, 1, 4, 2, 5]), [1, 2, 3, 4, 5])
    def test_duplicate_elements_list(self):
        self.assertEqual(sorter([3, 1, 4, 2, 2, 5, 1]), [1, 1, 2, 2, 3, 4, 5])
    def test_negative_numbers_list(self):
        self.assertEqual(sorter([-5, -2, -8, -1, -3]), [-8, -5, -3, -2, -1])
    def test_mixed_data_types_list(self):
        self.assertEqual(sorter(['apple', 2, 'banana', 1, 'cherry']), [1, 2, 'apple', 'banana', 'cherry'])
    def test_large_input_list(self):
        self.assertEqual(sorter(list(range(1000, 0, -1))), list(range(1, 1001)))
    def test_list_with_none_values(self):
        self.assertEqual(sorter([None, 2, None, 1, None]), [None, None, None, 1, 2])
    def test_list_with_nan_values(self):
        self.assertEqual(sorter([float('nan'), 2, float('nan'), 1, float('nan')]), [1, 2, float('nan'), float('nan'), float('nan')])
    def test_list_with_complex_numbers(self):
        self.assertEqual(sorter([3 + 2j, 1 + 1j, 4 + 3j, 2 + 1j, 5 + 4j]), [1 + 1j, 2 + 1j, 3 + 2j, 4 + 3j, 5 + 4j])
    def test_list_with_custom_class_objects(self):
        class Person:
            def __init__(self, name, age):
                self.name = name
                self.age = age
            def __repr__(self):
                return f\"Person('{self.name}', {self.age})\"
        input_list = [Person('Alice', 25), Person('Bob', 30), Person('Charlie', 20)]
        expected_output = [Person('Charlie', 20), Person('Alice', 25), Person('Bob', 30)]
        self.assertEqual(sorter(input_list), expected_output)
    def test_list_with_uncomparable_elements(self):
        with self.assertRaises(TypeError):
            sorter([5, 'apple', 3, [1, 2, 3], 2])
    def test_list_with_custom_comparison_function(self):
        input_list = [5, 4, 3, 2, 1]
        expected_output = [5, 4, 3, 2, 1]
        self.assertEqual(sorter(input_list, reverse=True), expected_output)
if __name__ == '__main__':
    unittest.main()
"""

    expected_remove_asserts = """from code_to_optimize.bubble_sort import sorter
import unittest
def sorter(arr):
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp
    return arr
class SorterTestCase(unittest.TestCase):
    def test_empty_list(self):
        codeflash_output = sorter([])
    def test_single_element_list(self):
        codeflash_output = sorter([5])
    def test_ascending_order_list(self):
        codeflash_output = sorter([1, 2, 3, 4, 5])
    def test_descending_order_list(self):
        codeflash_output = sorter([5, 4, 3, 2, 1])
    def test_random_order_list(self):
        codeflash_output = sorter([3, 1, 4, 2, 5])
    def test_duplicate_elements_list(self):
        codeflash_output = sorter([3, 1, 4, 2, 2, 5, 1])
    def test_negative_numbers_list(self):
        codeflash_output = sorter([-5, -2, -8, -1, -3])
    def test_mixed_data_types_list(self):
        codeflash_output = sorter(['apple', 2, 'banana', 1, 'cherry'])
    def test_large_input_list(self):
        codeflash_output = sorter(list(range(1000, 0, -1)))
    def test_list_with_none_values(self):
        codeflash_output = sorter([None, 2, None, 1, None])
    def test_list_with_nan_values(self):
        codeflash_output = sorter([float('nan'), 2, float('nan'), 1, float('nan')])
    def test_list_with_complex_numbers(self):
        codeflash_output = sorter([3 + 2j, 1 + 1j, 4 + 3j, 2 + 1j, 5 + 4j])
    def test_list_with_custom_class_objects(self):
        class Person:
            def __init__(self, name, age):
                self.name = name
                self.age = age
            def __repr__(self):
                return f\"Person('{self.name}', {self.age})\"
        input_list = [Person('Alice', 25), Person('Bob', 30), Person('Charlie', 20)]
        expected_output = [Person('Charlie', 20), Person('Alice', 25), Person('Bob', 30)]
        codeflash_output = sorter(input_list)
    def test_list_with_uncomparable_elements(self):
        with self.assertRaises(TypeError):
            sorter([5, 'apple', 3, [1, 2, 3], 2])
    def test_list_with_custom_comparison_function(self):
        input_list = [5, 4, 3, 2, 1]
        expected_output = [5, 4, 3, 2, 1]
        codeflash_output = sorter(input_list, reverse=True)
if __name__ == '__main__':
    unittest.main()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
"""

    expected_instrumentation = """import unittest

class SorterTestCase(unittest.TestCase):

    def test_empty_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_empty_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_single_element_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([5])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_single_element_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_ascending_order_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([1, 2, 3, 4, 5])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_ascending_order_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_descending_order_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([5, 4, 3, 2, 1])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_descending_order_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_random_order_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([3, 1, 4, 2, 5])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_random_order_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_duplicate_elements_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([3, 1, 4, 2, 2, 5, 1])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_duplicate_elements_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_negative_numbers_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([-5, -2, -8, -1, -3])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_negative_numbers_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_mixed_data_types_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind(['apple', 2, 'banana', 1, 'cherry'])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_mixed_data_types_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_large_input_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind(list(range(1000, 0, -1)))
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_large_input_list', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_list_with_none_values(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([None, 2, None, 1, None])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_none_values', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_list_with_nan_values(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([float('nan'), 2, float('nan'), 1, float('nan')])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_nan_values', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_list_with_complex_numbers(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(sorter).bind([3 + 2j, 1 + 1j, 4 + 3j, 2 + 1j, 5 + 4j])
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_complex_numbers', 'sorter', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_list_with_custom_class_objects(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])

        class Person:

            def __init__(self, name, age):
                self.name = name
                self.age = age

            def __repr__(self):
                return f"Person('{self.name}', {self.age})"
        input_list = [Person('Alice', 25), Person('Bob', 30), Person('Charlie', 20)]
        expected_output = [Person('Charlie', 20), Person('Alice', 25), Person('Bob', 30)]
        _call__bound__arguments = inspect.signature(sorter).bind(input_list)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_custom_class_objects', 'sorter', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value

    def test_list_with_uncomparable_elements(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        with self.assertRaises(TypeError):
            _call__bound__arguments = inspect.signature(sorter).bind([5, 'apple', 3, [1, 2, 3], 2])
            _call__bound__arguments.apply_defaults()
            codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_uncomparable_elements', 'sorter', '0_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

    def test_list_with_custom_comparison_function(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        input_list = [5, 4, 3, 2, 1]
        expected_output = [5, 4, 3, 2, 1]
        _call__bound__arguments = inspect.signature(sorter).bind(input_list, reverse=True)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(sorter, 'code_to_optimize_path', 'SorterTestCase', 'test_list_with_custom_comparison_function', 'sorter', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
if __name__ == '__main__':
    unittest.main()"""

    auxillary_functions: list[str] = []
    module_node = ast.parse(expected_remove_asserts)
    test_module_path = "code_to_optimize_path"
    function_to_optimize = FunctionToOptimize(function_name="sorter", file_path="/tmp/path", parents=[])
    remove_asserts = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert remove_asserts == expected_remove_asserts

    new_module_node = InjectPerfAndLogging(function_to_optimize, auxillary_functions, test_module_path).visit(
        module_node
    )
    assert ast.unparse(new_module_node).strip("\n") == expected_instrumentation


def test_InjectPerfAndLogging_preserves_nested_functions() -> None:
    """Test that nested functions containing target function calls are preserved.

    This is a regression test for a bug where nested functions (like thread workers)
    that called the target function were being deleted and replaced with the
    instrumented call at the test function level.
    """
    code = """
def test_concurrent_calls():
    agent = MyClass()
    results = [None] * 10

    def worker(idx):
        results[idx] = agent.target_method()

    threads = [threading.Thread(target=worker, args=(i,)) for i in range(10)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
"""
    auxillary_functions: list[str] = []
    module_node = ast.parse(code)
    test_module_path = "test_path"
    function_to_optimize = FunctionToOptimize(
        function_name="target_method", file_path="/tmp/path", parents=[FunctionParent(name="MyClass", type="ClassDef")]
    )
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.BEHAVIOR
    ).visit(module_node)

    result = ast.unparse(new_module_node)

    # The nested worker function must be preserved
    assert "def worker(idx):" in result
    # The body of the nested function should still contain the original call
    assert "results[idx] = agent.target_method()" in result
    # The function should NOT be replaced with codeflash_wrap at the test level
    # (the call inside the nested function is not instrumented, which is expected)


def test_InjectPerfAndLogging_preserves_nested_async_functions() -> None:
    """Test that nested async functions containing target function calls are preserved."""
    code = """
def test_async_concurrent():
    agent = MyClass()

    async def async_worker():
        return agent.target_method()

    result = asyncio.run(async_worker())
"""
    auxillary_functions: list[str] = []
    module_node = ast.parse(code)
    test_module_path = "test_path"
    function_to_optimize = FunctionToOptimize(
        function_name="target_method", file_path="/tmp/path", parents=[FunctionParent(name="MyClass", type="ClassDef")]
    )
    new_module_node = InjectPerfAndLogging(
        function_to_optimize, auxillary_functions, test_module_path, mode=TestingMode.BEHAVIOR
    ).visit(module_node)

    result = ast.unparse(new_module_node)

    # The nested async function must be preserved
    assert "async def async_worker():" in result
    # The body of the nested function should still contain the original call
    assert "return agent.target_method()" in result
