"""Exhaustive unit tests for instrument_test_source with different used_frameworks values.

These tests verify that the wrapper function is correctly generated with GPU device
synchronization code for different framework imports (torch, tensorflow, jax).
"""

from __future__ import annotations

from aiservice.models.functions_to_optimize import FunctionToOptimize
from core.languages.python.testgen.instrumentation.instrument_new_tests import instrument_test_source
from core.languages.python.testgen.models import TestingMode
from tests.conftest import normalize_code

# ============================================================================
# Expected wrapper function strings for BEHAVIOR mode
# ============================================================================

# Baseline: No GPU frameworks
CODEFLASH_WRAP_NO_FRAMEWORKS_BEHAVIOR = """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"""


# PyTorch only (import torch)
CODEFLASH_WRAP_TORCH_BEHAVIOR = """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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        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"""


# PyTorch with alias (import torch as th)
CODEFLASH_WRAP_TORCH_ALIASED_BEHAVIOR = """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
    _codeflash_should_sync_cuda = th.cuda.is_available() and th.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(th.backends, 'mps') and th.backends.mps.is_available() and hasattr(th.mps, 'synchronize')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            th.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            th.mps.synchronize()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            th.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            th.mps.synchronize()
        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"""


# TensorFlow only (import tensorflow)
CODEFLASH_WRAP_TENSORFLOW_BEHAVIOR = """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
    _codeflash_should_sync_tf = hasattr(tensorflow.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        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"""


# TensorFlow with alias (import tensorflow as tf)
CODEFLASH_WRAP_TENSORFLOW_ALIASED_BEHAVIOR = """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
    _codeflash_should_sync_tf = hasattr(tf.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_tf:
            tf.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_tf:
            tf.test.experimental.sync_devices()
        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"""


# JAX only (import jax)
CODEFLASH_WRAP_JAX_BEHAVIOR = """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
    _codeflash_should_sync_jax = hasattr(jax, 'block_until_ready')
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_jax:
            jax.block_until_ready(return_value)
        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"""


# JAX with alias (import jax as jnp)
CODEFLASH_WRAP_JAX_ALIASED_BEHAVIOR = """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
    _codeflash_should_sync_jax = hasattr(jnp, 'block_until_ready')
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_jax:
            jnp.block_until_ready(return_value)
        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"""


# Torch + TensorFlow
CODEFLASH_WRAP_TORCH_TENSORFLOW_BEHAVIOR = """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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    _codeflash_should_sync_tf = hasattr(tensorflow.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        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"""


# All three frameworks: torch, tensorflow, jax
CODEFLASH_WRAP_ALL_FRAMEWORKS_BEHAVIOR = """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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    _codeflash_should_sync_jax = hasattr(jax, 'block_until_ready')
    _codeflash_should_sync_tf = hasattr(tensorflow.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_jax:
            jax.block_until_ready(return_value)
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        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"""


# ============================================================================
# Expected wrapper function strings for PERFORMANCE mode
# ============================================================================

CODEFLASH_WRAP_NO_FRAMEWORKS_PERFORMANCE = """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"""


CODEFLASH_WRAP_TORCH_PERFORMANCE = """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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        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"""


CODEFLASH_WRAP_JAX_PERFORMANCE = """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
    _codeflash_should_sync_jax = hasattr(jax, 'block_until_ready')
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_jax:
            jax.block_until_ready(return_value)
        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"""


CODEFLASH_WRAP_TENSORFLOW_PERFORMANCE = """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
    _codeflash_should_sync_tf = hasattr(tensorflow.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        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"""


CODEFLASH_WRAP_ALL_FRAMEWORKS_PERFORMANCE = """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
    _codeflash_should_sync_cuda = torch.cuda.is_available() and torch.cuda.is_initialized()
    _codeflash_should_sync_mps = not _codeflash_should_sync_cuda and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() and hasattr(torch.mps, 'synchronize')
    _codeflash_should_sync_jax = hasattr(jax, 'block_until_ready')
    _codeflash_should_sync_tf = hasattr(tensorflow.test.experimental, 'sync_devices')
    gc.disable()
    try:
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        if _codeflash_should_sync_cuda:
            torch.cuda.synchronize()
        elif _codeflash_should_sync_mps:
            torch.mps.synchronize()
        if _codeflash_should_sync_jax:
            jax.block_until_ready(return_value)
        if _codeflash_should_sync_tf:
            tensorflow.test.experimental.sync_devices()
        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"""


# ============================================================================
# Helper functions
# ============================================================================


def build_expected_behavior_test(imports: str, wrapper: str, test_body: str) -> str:
    """Build expected output for BEHAVIOR mode tests."""
    return f"""{imports}


{wrapper}

def test_my_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
{test_body}"""


def build_expected_performance_test(imports: str, wrapper: str, test_body: str) -> str:
    """Build expected output for PERFORMANCE mode tests."""
    return f"""{imports}


{wrapper}

def test_my_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
{test_body}"""


# ============================================================================
# BEHAVIOR MODE TESTS
# ============================================================================


class TestInstrumentTestSourceNoFrameworksBehavior:
    """Tests for instrument_test_source with no GPU frameworks in BEHAVIOR mode."""

    def test_no_frameworks_behavior_mode(self) -> None:
        """Test instrumentation with no GPU framework imports."""
        code = """from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
from mymodule import my_function"""

        test_body = """    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = build_expected_behavior_test(imports, CODEFLASH_WRAP_NO_FRAMEWORKS_BEHAVIOR, test_body)

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceTorchBehavior:
    """Tests for instrument_test_source with PyTorch in BEHAVIOR mode."""

    def test_torch_import_behavior_mode(self) -> None:
        """Test instrumentation with PyTorch import."""
        code = """import torch
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import torch
from mymodule import my_function"""

        # Note: torch.manual_seed(42) is automatically added for reproducibility
        test_body = """    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_torch_aliased_import_behavior_mode(self) -> None:
        """Test instrumentation with PyTorch imported as alias."""
        code = """import torch as th
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import torch as th
from mymodule import my_function"""

        # Note: th.manual_seed(42) is automatically added for reproducibility
        test_body = """    th.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_ALIASED_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_torch_submodule_import_behavior_mode(self) -> None:
        """Test instrumentation with PyTorch submodule import (from torch import nn)."""
        code = """from torch import nn
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        # Note: 'import torch' is automatically added for GPU sync and seeding
        imports = """from __future__ import annotations

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

import dill as pickle
import torch
from mymodule import my_function
from torch import nn"""

        # Note: torch.manual_seed(42) is automatically added for reproducibility
        test_body = """    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceTensorFlowBehavior:
    """Tests for instrument_test_source with TensorFlow in BEHAVIOR mode."""

    def test_tensorflow_import_behavior_mode(self) -> None:
        """Test instrumentation with TensorFlow import."""
        code = """import tensorflow
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import tensorflow
from mymodule import my_function"""

        # Note: tensorflow.random.set_seed(42) is automatically added for reproducibility
        test_body = """    tensorflow.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TENSORFLOW_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_tensorflow_aliased_import_behavior_mode(self) -> None:
        """Test instrumentation with TensorFlow imported as alias."""
        code = """import tensorflow as tf
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import tensorflow as tf
from mymodule import my_function"""

        # Note: tf.random.set_seed(42) is automatically added for reproducibility
        test_body = """    tf.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TENSORFLOW_ALIASED_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceJaxBehavior:
    """Tests for instrument_test_source with JAX in BEHAVIOR mode."""

    def test_jax_import_behavior_mode(self) -> None:
        """Test instrumentation with JAX import."""
        code = """import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import jax
from mymodule import my_function"""

        test_body = """    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = build_expected_behavior_test(imports, CODEFLASH_WRAP_JAX_BEHAVIOR, test_body)

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_jax_aliased_import_behavior_mode(self) -> None:
        """Test instrumentation with JAX imported as alias."""
        code = """import jax as jnp
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import jax as jnp
from mymodule import my_function"""

        test_body = """    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = build_expected_behavior_test(imports, CODEFLASH_WRAP_JAX_ALIASED_BEHAVIOR, test_body)

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceMultipleFrameworksBehavior:
    """Tests for instrument_test_source with multiple frameworks in BEHAVIOR mode."""

    def test_torch_and_tensorflow_behavior_mode(self) -> None:
        """Test instrumentation with both PyTorch and TensorFlow imports."""
        code = """import torch
import tensorflow
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import tensorflow
import torch
from mymodule import my_function"""

        # Note: Both torch.manual_seed(42) and tensorflow.random.set_seed(42) are added for reproducibility
        test_body = """    torch.manual_seed(42)
    tensorflow.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_TENSORFLOW_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_all_three_frameworks_behavior_mode(self) -> None:
        """Test instrumentation with PyTorch, TensorFlow, and JAX imports."""
        code = """import torch
import tensorflow
import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

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

import dill as pickle
import jax
import tensorflow
import torch
from mymodule import my_function"""

        # Note: Both torch.manual_seed(42) and tensorflow.random.set_seed(42) are added (JAX doesn't have a global seed)
        test_body = """    torch.manual_seed(42)
    tensorflow.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_ALL_FRAMEWORKS_BEHAVIOR}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


# ============================================================================
# PERFORMANCE MODE TESTS
# ============================================================================


class TestInstrumentTestSourceNoFrameworksPerformance:
    """Tests for instrument_test_source with no GPU frameworks in PERFORMANCE mode."""

    def test_no_frameworks_performance_mode(self) -> None:
        """Test instrumentation with no GPU framework imports in PERFORMANCE mode."""
        code = """from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

import gc
import os
import time

from mymodule import my_function"""

        test_body = """    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, 1, 2)
    result = codeflash_return_value
    assert result == 3"""

        expected = build_expected_performance_test(imports, CODEFLASH_WRAP_NO_FRAMEWORKS_PERFORMANCE, test_body)

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.PERFORMANCE,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceTorchPerformance:
    """Tests for instrument_test_source with PyTorch in PERFORMANCE mode."""

    def test_torch_import_performance_mode(self) -> None:
        """Test instrumentation with PyTorch import in PERFORMANCE mode."""
        code = """import torch
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

import gc
import os
import time

import torch
from mymodule import my_function"""

        # Note: torch.manual_seed(42) is automatically added for reproducibility
        test_body = """    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, 1, 2)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_PERFORMANCE}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.PERFORMANCE,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceJaxPerformance:
    """Tests for instrument_test_source with JAX in PERFORMANCE mode."""

    def test_jax_import_performance_mode(self) -> None:
        """Test instrumentation with JAX import in PERFORMANCE mode."""
        code = """import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

import gc
import os
import time

import jax
from mymodule import my_function"""

        test_body = """    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, 1, 2)
    result = codeflash_return_value
    assert result == 3"""

        expected = build_expected_performance_test(imports, CODEFLASH_WRAP_JAX_PERFORMANCE, test_body)

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.PERFORMANCE,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceTensorFlowPerformance:
    """Tests for instrument_test_source with TensorFlow in PERFORMANCE mode."""

    def test_tensorflow_import_performance_mode(self) -> None:
        """Test instrumentation with TensorFlow import in PERFORMANCE mode."""
        code = """import tensorflow
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

import gc
import os
import time

import tensorflow
from mymodule import my_function"""

        # Note: tensorflow.random.set_seed(42) is automatically added for reproducibility
        test_body = """    tensorflow.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, 1, 2)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TENSORFLOW_PERFORMANCE}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.PERFORMANCE,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


class TestInstrumentTestSourceAllFrameworksPerformance:
    """Tests for instrument_test_source with all frameworks in PERFORMANCE mode."""

    def test_all_frameworks_performance_mode(self) -> None:
        """Test instrumentation with PyTorch, TensorFlow, and JAX imports in PERFORMANCE mode."""
        code = """import torch
import tensorflow
import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        imports = """from __future__ import annotations

import gc
import os
import time

import jax
import tensorflow
import torch
from mymodule import my_function"""

        # Note: Both torch.manual_seed(42) and tensorflow.random.set_seed(42) are added (JAX doesn't have a global seed)
        test_body = """    torch.manual_seed(42)
    tensorflow.random.set_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, 1, 2)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_ALL_FRAMEWORKS_PERFORMANCE}

def test_my_function():
{test_body}"""

        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.PERFORMANCE,
        )

        assert result is not None
        assert normalize_code(result) == normalize_code(expected)


# ============================================================================
# EDGE CASES
# ============================================================================


class TestInstrumentTestSourceEdgeCases:
    """Edge case tests for instrument_test_source."""

    def test_torch_dotted_import_behavior_mode(self) -> None:
        """Test instrumentation with torch.cuda or torch.nn import."""
        code = """import torch.cuda
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        imports = """from __future__ import annotations

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

import dill as pickle
import torch
import torch.cuda
from mymodule import my_function"""

        test_body = """def test_my_function():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_my_function', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result = codeflash_return_value
    assert result == 3"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_BEHAVIOR}

{test_body}"""

        # Should detect torch from "torch.cuda" import
        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_multiple_function_calls_behavior_mode(self) -> None:
        """Test instrumentation with multiple function calls in the same test."""
        code = """import torch
from mymodule import my_function

def test_multiple_calls():
    result1 = my_function(1, 2)
    result2 = my_function(3, 4)
    assert result1 == 3
    assert result2 == 7
"""
        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        imports = """from __future__ import annotations

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

import dill as pickle
import torch
from mymodule import my_function"""

        test_body = """def test_multiple_calls():
    torch.manual_seed(42)
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_multiple_calls', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result1 = codeflash_return_value
    _call__bound__arguments = inspect.signature(my_function).bind(3, 4)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(my_function, 'test_module', None, 'test_multiple_calls', 'my_function', '1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    result2 = codeflash_return_value
    assert result1 == 3
    assert result2 == 7"""

        expected = f"""{imports}


{CODEFLASH_WRAP_TORCH_BEHAVIOR}

{test_body}"""
        assert result is not None
        assert normalize_code(result) == normalize_code(expected)

    def test_unittest_class_with_jax_behavior_mode(self) -> None:
        """Test instrumentation of unittest.TestCase with JAX import."""
        code = """import unittest
import jax
from mymodule import my_function


class TestMyFunction(unittest.TestCase):
    def test_basic(self):
        result = my_function(1, 2)
        self.assertEqual(result, 3)
"""
        func = FunctionToOptimize(function_name="my_function", parents=[], file_path="mymodule.py")

        result = instrument_test_source(
            test_source=code,
            function_to_optimize=func,
            helper_function_names=[],
            module_path="mymodule",
            test_module_path="test_module",
            test_framework="pytest",
            test_timeout=15,
            mode=TestingMode.BEHAVIOR,
        )

        imports = """from __future__ import annotations

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

import dill as pickle
import jax
from mymodule import my_function"""

        test_body = """class TestMyFunction(unittest.TestCase):

    def test_basic(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        _call__bound__arguments = inspect.signature(my_function).bind(1, 2)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(my_function, 'test_module', 'TestMyFunction', 'test_basic', 'my_function', '0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        result = codeflash_return_value
        self.assertEqual(result, 3)"""

        expected = f"""{imports}


{CODEFLASH_WRAP_JAX_BEHAVIOR}

{test_body}"""
        assert result is not None
        assert normalize_code(result) == normalize_code(expected)
