from __future__ import annotations

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

import pytest
from inference.core.workflows.execution_engine.entities.types import Kind
from inference.core.workflows.execution_engine.introspection.entities import (
    ReferenceDefinition, SelectorDefinition)
from inference.core.workflows.execution_engine.introspection.schema_parser import \
    retrieve_selectors_from_union_definition

def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    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}######$!")
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f"!######{test_stdout_tag}:{codeflash_duration}######!")
    if exception:
        raise exception
    return return_value

def test_retrieve_selectors_from_union_definition_with_anyof_single_reference():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test basic functionality with a single anyOf union containing one reference.\n    This verifies that the function correctly extracts selectors from a simple anyOf union.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/SomeModel', 'x-selected-element': 'predictions'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_anyof_single_reference', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='test_property', property_description='A test property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_oneof_single_reference():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test basic functionality with a single oneOf union containing one reference.\n    Ensures oneOf is handled the same way as anyOf.\n    '
    union_definition = {'oneOf': [{'$ref': '#/components/schemas/AnotherModel', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_oneof_single_reference', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='my_selector', property_description='My selector description', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_allof_single_reference():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test basic functionality with a single allOf union containing one reference.\n    Ensures allOf is handled correctly.\n    '
    union_definition = {'allOf': [{'$ref': '#/components/schemas/CombinedModel', 'x-selected-element': 'combined_output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_allof_single_reference', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='combined', property_description='Combined property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_multiple_references_same_element():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test merging of references pointing to the same selected element.\n    This verifies that references with identical selected_element are properly merged.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'same_element'}, {'$ref': '#/components/schemas/Model2', 'x-selected-element': 'same_element'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_multiple_references_same_element', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='merged_property', property_description='Property with merged references', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_multiple_references_different_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of multiple references pointing to different selected elements.\n    Verifies that distinct references are preserved separately.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'element_a'}, {'$ref': '#/components/schemas/Model2', 'x-selected-element': 'element_b'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_multiple_references_different_elements', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='multi_element_property', property_description='Property with multiple elements', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output
    selected_elements = {ref.selected_element for ref in result.allowed_references}

def test_retrieve_selectors_from_union_definition_preserves_property_metadata():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that property metadata (dimensionality offset, description) is preserved in result.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_preserves_property_metadata', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='test_prop', property_description='Test description', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=2, is_dimensionality_reference_property=True, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_empty_union():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior when union_definition has empty arrays for all union types.\n    Should return None when no valid references are found.\n    '
    union_definition = {'anyOf': [], 'oneOf': [], 'allOf': []}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_empty_union', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='empty_union', property_description='Empty union', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_no_union_keys():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior when union_definition has no anyOf, oneOf, or allOf keys.\n    Should return None when there are no union types.\n    '
    union_definition = {}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_no_union_keys', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='no_union_keys', property_description='No union keys', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_non_reference_union_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that non-reference union members are filtered out (return None from recursive calls).\n    Only reference definitions should be included in final result.\n    '
    union_definition = {'anyOf': [{'type': 'string'}, {'$ref': '#/components/schemas/Model1', 'x-selected-element': 'valid_element'}, {'type': 'number'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_non_reference_union_types', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='mixed_union', property_description='Mixed union with non-references', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_all_non_references():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior when union contains only non-reference types.\n    Should return None when no references can be extracted.\n    '
    union_definition = {'anyOf': [{'type': 'string'}, {'type': 'number'}, {'type': 'boolean'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_all_non_references', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='non_reference_union', property_description='Union with no references', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_list_element_flag():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that is_list_element flag is preserved in the result.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_list_element_flag', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='list_property', property_description='List property', union_definition=union_definition, is_list_element=True, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_dict_element_flag():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that is_dict_element flag is preserved in the result.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_dict_element_flag', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='dict_property', property_description='Dict property', union_definition=union_definition, is_list_element=False, is_dict_element=True, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_combined_list_and_dict_flags():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that both list and dict flags are properly combined when present in union results.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_combined_list_and_dict_flags', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='combined_property', property_description='Combined flags property', union_definition=union_definition, is_list_element=True, is_dict_element=True, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_mixed_union_types():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test union_definition that contains multiple union type keys (anyOf, oneOf, allOf).\n    All should be processed and results merged.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'element_a'}], 'oneOf': [{'$ref': '#/components/schemas/Model2', 'x-selected-element': 'element_b'}], 'allOf': [{'$ref': '#/components/schemas/Model3', 'x-selected-element': 'element_c'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_mixed_union_types', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='mixed_types', property_description='Mixed union types', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output
    selected_elements = {ref.selected_element for ref in result.allowed_references}

def test_retrieve_selectors_from_union_definition_with_negative_dimensionality_offset():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of negative dimensionality_offset values.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_negative_dimensionality_offset', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='negative_offset', property_description='Negative offset property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=-1, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_empty_property_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior with empty string as property_name.\n    Function should still process it without error.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_empty_property_name', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='', property_description='Empty name property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_empty_property_description():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior with empty string as property_description.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_empty_property_description', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='test_prop', property_description='', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_very_long_property_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test behavior with very long property_name string.\n    '
    long_name = 'a' * 1000
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_very_long_property_name', 'retrieve_selectors_from_union_definition', '3', codeflash_loop_index, property_name=long_name, property_description='Long name property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_special_characters_in_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of special characters in property names.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    special_name = 'prop_with-special.chars$123'
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_special_characters_in_names', 'retrieve_selectors_from_union_definition', '3', codeflash_loop_index, property_name=special_name, property_description='Special chars', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_unicode_property_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of unicode characters in property names.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    unicode_name = 'プロパティ名_属性'
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_unicode_property_names', 'retrieve_selectors_from_union_definition', '3', codeflash_loop_index, property_name=unicode_name, property_description='Unicode property', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_duplicate_references():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling when same reference appears multiple times in union.\n    Duplicates with same selected_element should be merged.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'element_a'}, {'$ref': '#/components/schemas/Model1', 'x-selected-element': 'element_a'}, {'$ref': '#/components/schemas/Model1', 'x-selected-element': 'element_a'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_duplicate_references', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='duplicates', property_description='Duplicate references', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_many_references():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test performance with a large number of distinct references in union.\n    Verifies the function can handle 100 different references efficiently.\n    '
    union_types = []
    for i in range(100):
        union_types.append({'$ref': f'#/components/schemas/Model{i}', 'x-selected-element': f'element_{i}'})
    union_definition = {'anyOf': union_types}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_many_references', 'retrieve_selectors_from_union_definition', '4', codeflash_loop_index, property_name='many_refs', property_description='Union with many references', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output
    selected_elements = {ref.selected_element for ref in result.allowed_references}
    for i in range(100):
        pass

def test_retrieve_selectors_from_union_definition_with_many_duplicates():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of 500 references that point to only 10 distinct elements.\n    Verifies proper merging of duplicates at scale.\n    '
    union_types = []
    for i in range(500):
        element_index = i % 10
        union_types.append({'$ref': f'#/components/schemas/Model{i}', 'x-selected-element': f'element_{element_index}'})
    union_definition = {'anyOf': union_types}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_many_duplicates', 'retrieve_selectors_from_union_definition', '4', codeflash_loop_index, property_name='many_duplicates', property_description='Union with many duplicates', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output
    selected_elements = {ref.selected_element for ref in result.allowed_references}

def test_retrieve_selectors_from_union_definition_with_mixed_large_union():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test with large union containing mix of anyOf, oneOf, and allOf.\n    Creates 300 total references across different union type keys.\n    '
    union_definition = {'anyOf': [{'$ref': f'#/components/schemas/ModelAny{i}', 'x-selected-element': f'elem_any_{i % 20}'} for i in range(100)], 'oneOf': [{'$ref': f'#/components/schemas/ModelOne{i}', 'x-selected-element': f'elem_one_{i % 20}'} for i in range(100)], 'allOf': [{'$ref': f'#/components/schemas/ModelAll{i}', 'x-selected-element': f'elem_all_{i % 20}'} for i in range(100)]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_mixed_large_union', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='large_mixed', property_description='Large mixed union', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_performance_with_large_batches():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test performance with large batch parameters.\n    Verifies function handles large sets of batch-accepting inputs.\n    '
    large_batch_set = {f'input_{i}' for i in range(500)}
    large_batch_scalars_set = {f'input_scalar_{i}' for i in range(300)}
    large_auto_cast_set = {f'input_auto_{i}' for i in range(200)}
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_performance_with_large_batches', 'retrieve_selectors_from_union_definition', '5', codeflash_loop_index, property_name='large_batch_property', property_description='Property with large batch sets', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=large_batch_set, inputs_accepting_batches_and_scalars=large_batch_scalars_set, inputs_enforcing_auto_batch_casting=large_auto_cast_set)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_with_large_dimensionality_offset():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test handling of very large dimensionality_offset values.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'output'}]}
    codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_with_large_dimensionality_offset', 'retrieve_selectors_from_union_definition', '2', codeflash_loop_index, property_name='large_offset', property_description='Large dimensionality offset', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=999, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_retrieve_selectors_from_union_definition_consistency_with_repeated_calls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    '\n    Test that repeated calls with same parameters produce consistent results.\n    Ensures no side effects or state pollution.\n    '
    union_definition = {'anyOf': [{'$ref': '#/components/schemas/Model1', 'x-selected-element': 'elem_a'}, {'$ref': '#/components/schemas/Model2', 'x-selected-element': 'elem_b'}]}
    results = []
    for _ in range(10):
        codeflash_return_value = codeflash_wrap(retrieve_selectors_from_union_definition, 'tests.inference.unit_tests.test_retrieve_selectors_from_union_definition__unit_test_1', None, 'test_retrieve_selectors_from_union_definition_consistency_with_repeated_calls', 'retrieve_selectors_from_union_definition', '3_0', codeflash_loop_index, property_name='consistent', property_description='Consistency test', union_definition=union_definition, is_list_element=False, is_dict_element=False, property_dimensionality_offset=0, is_dimensionality_reference_property=False, inputs_accepting_batches=set(), inputs_accepting_batches_and_scalars=set(), inputs_enforcing_auto_batch_casting=set())
        codeflash_output = codeflash_return_value
        result = codeflash_output
        results.append(result)
    for r in results[1:]:
        pass