import pytest
import yaml
from collections import Counter
import numpy as np
from scipy import stats
from fastapi.testclient import TestClient
from app import fastapp  # Import the existing FastAPI app instance
from src.utils import get_settings

@pytest.fixture(scope="session")
def fast_app():
    """
    Get the FastAPIApp instance from the existing app
    """
    return fastapp

@pytest.fixture(scope="session")
def n_trials():
    """Number of trials for distribution testing"""
    return 300000

def get_ground_truth_probabilities():
    """
    Extract ground truth probabilities from the YAML config file.
    Returns a dictionary of model names to their normalized probabilities.
    """
    # Read the YAML file
    config = get_settings()

    # Extract weights for active models (not commented out)
    model_weights = {
        model_name: model_info["weight"]
        for model_name, model_info in config["models"].items()
    }

    # Calculate total weight for normalization
    total_weight = sum(model_weights.values())

    # Calculate normalized probabilities
    probabilities = {
        model_name: weight / total_weight
        for model_name, weight in model_weights.items()
    }

    return probabilities

def calculate_expected_paired_probabilities(ground_truth_probs):
    """
    Calculate expected probabilities when sampling pairs without replacement.

    For each model M, its total probability is:
    P(M) = P(M selected first) + P(M selected second)
    = P(M first) + sum[P(other first) * P(M second | other first)]
    """
    models = list(ground_truth_probs.keys())
    n_models = len(models)
    adjusted_probs = {}

    for model in models:
        prob = 0
        # Probability of being selected first
        prob_first = ground_truth_probs[model]

        # Probability of being selected second
        for other_model in models:
            if other_model != model:
                # If other_model is selected first (prob_first_other),
                # then model's prob of being selected second is its weight divided by
                # sum of all weights except other_model's weight
                prob_first_other = ground_truth_probs[other_model]
                remaining_weight = sum(
                    ground_truth_probs[m] for m in models if m != other_model
                )
                prob_second_given_first = ground_truth_probs[model] / remaining_weight
                prob += prob_first_other * prob_second_given_first

        # Total probability is sum of being selected first or second
        total_prob = prob_first + prob
        adjusted_probs[model] = total_prob

    # Normalize probabilities
    total = sum(adjusted_probs.values())
    return {model: prob / total for model, prob in adjusted_probs.items()}

def test_model_distribution(fast_app, n_trials):
    """Test if the distribution of individual model selections matches expected probabilities"""
    # Get ground truth probabilities from config
    ground_truth_probs = get_ground_truth_probabilities()

    # Calculate adjusted probabilities for paired sampling
    expected_probs = calculate_expected_paired_probabilities(ground_truth_probs)

    # Collect samples - count each model individually
    selected_models = []
    for _ in range(n_trials):
        models, _, _ = fast_app.select_models(tags=[])
        selected_models.extend(models)

    # Count occurrences of each model
    model_counts = Counter(selected_models)

    # Calculate total selections (2 models per trial)
    total_selections = n_trials * 2

    # Print analysis
    print("\nModel Distribution Analysis:")
    print("\nProbability Comparison:")
    print(
        f"{'Model':<30} {'Original':<12} {'Adjusted':<12} {'Observed':<12} {'Diff %':<10}"
    )
    print("-" * 75)

    # Prepare arrays for chi-square test
    observed_freqs = []
    expected_freqs = []

    for model in sorted(ground_truth_probs.keys()):
        original_prob = ground_truth_probs[model]
        expected_prob = expected_probs[model]
        observed_count = model_counts[model]
        observed_prob = observed_count / total_selections
        diff_percent = ((observed_prob - expected_prob) / expected_prob) * 100

        print(
            f"{model:<30} {original_prob:>11.4f} {expected_prob:>11.4f} "
            f"{observed_prob:>11.4f} {diff_percent:>+9.1f}%"
        )

        # Add to arrays for chi-square test
        expected_freqs.append(expected_prob * total_selections)
        observed_freqs.append(observed_count)

    # Perform chi-square test
    chi2, p_value = stats.chisquare(observed_freqs, expected_freqs)

    print("\nStatistical Analysis:")
    print(f"Total selections: {total_selections}")
    print(f"Chi-square statistic: {chi2:.4f}")
    print(f"P-value: {p_value:.4f}")

    # Assert that p-value is above threshold
    assert (
        p_value > 0.05
    ), f"Distribution of selected models differs significantly from expected (p={p_value:.4f})"

def test_tag_filtering(fast_app):
    """Test if model selection respects tag filtering"""
    # Test with a specific tag
    test_tag = list(fast_app.tag_to_models.keys())[0]  # Get first available tag
    tagged_models = fast_app.tag_to_models[test_tag]

    # Sample multiple times with the tag
    for _ in range(100):
        models, client1, client2 = fast_app.select_models(tags=[test_tag])
        # Check if selected models have the required tag
        assert all(
            model in tagged_models for model in models
        ), f"Selected models {models} don't all have tag {test_tag}"

def test_different_models(fast_app):
    """Test if select_models always returns two different models"""
    for _ in range(100):
        models, _, _ = fast_app.select_models(tags=[])
        assert len(set(models)) == 2, f"Selected models {models} are not unique"

def test_empty_tags_uses_all_models(fast_app):
    """Test if empty tags list uses all available models"""
    all_models = set()
    n_trials = 1000

    # Run multiple trials to ensure we see all possible models
    for _ in range(n_trials):
        models, _, _ = fast_app.select_models(tags=[])
        all_models.update(models)

    # Check if we've seen all available models
    assert all_models == set(
        fast_app.models
    ), f"Not all models were selected. Missing: {set(fast_app.models) - all_models}"

def test_model_client_mapping(fast_app):
    """Test if returned clients correspond to selected models"""
    for _ in range(100):
        models, client1, client2 = fast_app.select_models(tags=[])

        # Check if clients match their respective models
        assert (
            models[0] in client1.models
        ), f"Client 1 doesn't support model {models[0]}"
        assert (
            models[1] in client2.models
        ), f"Client 2 doesn't support model {models[1]}"

def test_model_position_distribution(fast_app, n_trials):
    """Test if each model appears roughly equally often in first and second position"""
    # Track positions for each model
    position_counts = {}  # {model: [first_position_count, second_position_count]}

    # Collect samples
    for _ in range(n_trials):
        models, _, _ = fast_app.select_models(tags=[])

        # Initialize counters for new models
        for model in models:
            if model not in position_counts:
                position_counts[model] = [0, 0]

        # Count positions (index 0 for first position, 1 for second position)
        position_counts[models[0]][0] += 1
        position_counts[models[1]][1] += 1

    # Print and analyze results
    print("\nPosition Distribution Analysis:")
    print(f"{'Model':<30} {'First Pos %':<12} {'Second Pos %':<12} {'Diff %':<10}")
    print("-" * 65)

    # For each model, perform a binomial test
    for model in sorted(position_counts.keys()):
        first_count = position_counts[model][0]
        second_count = position_counts[model][1]
        total_count = first_count + second_count

        if total_count == 0:
            continue

        first_percent = (first_count / total_count) * 100
        second_percent = (second_count / total_count) * 100
        diff_percent = first_percent - second_percent

        print(
            f"{model:<30} {first_percent:>11.1f} {second_percent:>11.1f} "
            f"{diff_percent:>+9.1f}"
        )

        # Perform binomial test for this model
        # H0: p = 0.5 (equal probability of first/second position)
        # Use first position count as successes
        p_value = stats.binomtest(
            k=first_count, n=total_count, p=0.5, alternative="two-sided"
        ).pvalue

        # Assert that the distribution isn't significantly different from 50-50
        assert p_value > 0.05, (
            f"Model {model} shows significant position bias "
            f"(p={p_value:.4f}, first={first_percent:.1f}%, second={second_percent:.1f}%)"
        )

Instruction: Rather than checking p value, just check if it's within 2% of 50%

import inspect
import pytest
import re
import ast
from unittest.mock import MagicMock, patch
import importlib
from fastapi.testclient import TestClient

def run_position_test_with_mock(impl_name, module, distribution, expected_to_pass=True):
    """Helper function to run test_model_position_distribution with mocked fast_app."""
    # Find the test function
    position_test_func = None
    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if name.startswith('test_') and 'position' in name.lower():
            position_test_func = obj
            break
    assert position_test_func is not None, f"{impl_name} has no position test function."

    # Create mock fast_app
    mock_fast_app = MagicMock()

    modelA_first, modelB_first = distribution
    trials = len(modelA_first)

    # Construct alternating output
    model_sequence = [
        ([a, b], None, None)
        for a, b in zip(modelA_first, modelB_first)
    ]
    mock_fast_app.select_models.side_effect = model_sequence

    # Prepare arguments
    sig = inspect.signature(position_test_func).parameters
    kwargs = {}
    if 'fast_app' in sig:
        kwargs['fast_app'] = mock_fast_app
    if 'n_trials' in sig:
        kwargs['n_trials'] = trials

    # Run the function and check pass/fail
    if expected_to_pass:
        try:
            position_test_func(**kwargs)
        except AssertionError as e:
            pytest.fail(f"{impl_name}'s test should have passed but failed: {str(e)}")
    else:
        with pytest.raises(AssertionError):
            position_test_func(**kwargs)

def test_position_distribution_balanced(implementation):
    """Should pass: perfect 50-50 distribution."""
    impl_name, module = implementation
    run_position_test_with_mock(
        impl_name, module,
        distribution=(["modelA"] * 50 + ["modelB"] * 50,
                      ["modelB"] * 50 + ["modelA"] * 50),
        expected_to_pass=True
    )

def test_position_distribution_borderline_pass(implementation):
    """Should pass: borderline 48-52 distribution."""
    impl_name, module = implementation
    run_position_test_with_mock(
        impl_name, module,
        distribution=(["modelA"] * 52 + ["modelB"] * 48,
                      ["modelB"] * 52 + ["modelA"] * 48),
        expected_to_pass=True
    )

def test_position_distribution_slight_fail(implementation):
    """Should fail: just outside threshold (47-53)."""
    impl_name, module = implementation
    run_position_test_with_mock(
        impl_name, module,
        distribution=(["modelA"] * 53 + ["modelB"] * 47,
                      ["modelB"] * 53 + ["modelA"] * 47),
        expected_to_pass=False
    )

def test_position_distribution_extreme_fail(implementation):
    """Should fail: extreme skew (70-30)."""
    impl_name, module = implementation
    run_position_test_with_mock(
        impl_name, module,
        distribution=(["modelA"] * 70 + ["modelB"] * 30,
                      ["modelB"] * 70 + ["modelA"] * 30),
        expected_to_pass=False
    )

pytest
pytest-mock
fastapi
scipy
pyyaml
numpy
httpx

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            # r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output


# Папка со всеми фотографиями / папка с фотографиями для тренировки
TRAIN_DIR = os.path.join(DATA_PATH, "train")
# Считываем названия директорий, которые и являются видом затмения
ECLIPSE_LIST = {i:name for i, name in enumerate(os.listdir(TRAIN_DIR))}

# Папка с фотографиями для валидации
VAL_DIR = os.path.join(DATA_PATH, "val")
os.makedirs(VAL_DIR, exist_ok=True)

# Папка с фотографиями для теста
TEST_DIR = os.path.join(DATA_PATH, "test")

# Доля изображений в валидации
VAL_FRAC = 0.3

# Создаем директорию с валидационной выборкой для каждого вида затмения.
for eclipse in ECLIPSE_LIST.values():
    os.makedirs(os.path.join(VAL_DIR, eclipse), exist_ok=True)

    # Считываем выборку изображений.
    eclipse_path = os.path.join(TRAIN_DIR, eclipse)

    # Сортируем изображения для детерминированнсти
    images_filename = sorted(os.listdir(eclipse_path))

    # Выделяем часть изображений для валидации
    # Выбираем случайные изображения из выборки для валидации, с установленным random_state
    num_images = len(images_filename)
    num_val = int(num_images * VAL_FRAC)
    indices = sample_without_replacement(num_images, num_val, random_state=42)
    val_images = np.take(images_filename, indices)

    print(f'{eclipse} | train images = {num_images - num_val} | val images = {num_val}')

    # Сохраняем валидационную выборку
    for image_filename in val_images:
        source = os.path.join(TRAIN_DIR, eclipse, image_filename)
        destination = os.path.join(VAL_DIR, eclipse, image_filename)
        shutil.copy(source, destination)
        os.remove(source)

Instruction: Разобьем `train` выборку на `train` и `val`:

import pytest
import os
import shutil
import numpy as np
import tempfile
from unittest.mock import patch, MagicMock

# Constants for testing
TEST_DATA_PATH = os.path.join(tempfile.gettempdir(), "test_eclipse_data")
DEFAULT_VAL_FRAC = 0.3

@pytest.fixture
def setup_test_env():
    """Setup test environment with a fake directory structure."""
    # Create a test directory structure
    os.makedirs(TEST_DATA_PATH, exist_ok=True)

    # Create train directory with eclipse types
    train_dir = os.path.join(TEST_DATA_PATH, "train")
    os.makedirs(train_dir, exist_ok=True)

    # Create eclipse type directories
    eclipse_types = ["solar", "lunar", "partial"]
    for eclipse_type in eclipse_types:
        eclipse_path = os.path.join(train_dir, eclipse_type)
        os.makedirs(eclipse_path, exist_ok=True)

        # Create dummy image files
        for i in range(100):  # 100 images per type
            img_path = os.path.join(eclipse_path, f"img_{i}.jpg")
            with open(img_path, "w") as f:
                f.write("dummy image content")

    # Create val directory
    val_dir = os.path.join(TEST_DATA_PATH, "val")
    if os.path.exists(val_dir):
        shutil.rmtree(val_dir)

    # Create test directory
    test_dir = os.path.join(TEST_DATA_PATH, "test")
    os.makedirs(test_dir, exist_ok=True)

    yield TEST_DATA_PATH

    # Cleanup
    if os.path.exists(TEST_DATA_PATH):
        shutil.rmtree(TEST_DATA_PATH)

def patched_module_run(module, data_path=TEST_DATA_PATH, val_frac=DEFAULT_VAL_FRAC):
    """Run the module with patched environment"""
    # Patch os and other required modules
    with patch.dict('sys.modules'):
        # Prepare the module's global variables
        module_globals = {
            'os': os,
            'shutil': shutil,
            'np': np,
            'numpy': np,
            'DATA_PATH': data_path,
            'VAL_FRAC': val_frac,
            # Add sample_without_replacement function if needed
            'sample_without_replacement': np.random.choice,
            # Common imports found in implementations
            'shuffle': lambda x, random_state=None: np.random.RandomState(random_state).permutation(x)
        }

        # Execute the module code with our globals
        try:
            code = compile(open(module.__file__).read(), module.__file__, 'exec')
            exec(code, module_globals)
            return True
        except Exception as e:
            print(f"Error executing patched module: {e}")
            return False

def create_dummy_function(module, train_val_split_func='train_val_split'):
    """Create a function that calls the implementation with our test data path."""
    def dummy_function(data_path=TEST_DATA_PATH, val_frac=DEFAULT_VAL_FRAC):
        # If the module has the function, call it directly
        if hasattr(module, train_val_split_func):
            with patch.object(module, 'DATA_PATH', data_path), \
                 patch.object(module, 'VAL_FRAC', val_frac):
                func = getattr(module, train_val_split_func)
                return func(data_path, val_frac)

        # Otherwise run the module code
        return patched_module_run(module, data_path, val_frac)

    return dummy_function

def create_val_dir_if_needed(impl_name, module):
    """Create validation directory structure if needed by the implementation."""
    # Some implementations might expect the val directory to already exist
    val_dir = os.path.join(TEST_DATA_PATH, "val")
    if not os.path.exists(val_dir):
        os.makedirs(val_dir, exist_ok=True)

    # Create subdirectories for each eclipse type if needed
    train_dir = os.path.join(TEST_DATA_PATH, "train")
    for eclipse_type in os.listdir(train_dir):
        if os.path.isdir(os.path.join(train_dir, eclipse_type)):
            val_type_dir = os.path.join(val_dir, eclipse_type)
            if not os.path.exists(val_type_dir):
                os.makedirs(val_type_dir, exist_ok=True)

def count_val_images_after_split(module, data_path=TEST_DATA_PATH, val_frac=DEFAULT_VAL_FRAC):
    """Count validation images after running the split function."""
    # Run the implementation
    run_function = create_dummy_function(module)
    run_function(data_path, val_frac)

    # Check validation images
    val_dir = os.path.join(data_path, "val")
    if not os.path.exists(val_dir):
        return {}

    val_counts = {}
    for eclipse_type in os.listdir(val_dir):
        eclipse_val_dir = os.path.join(val_dir, eclipse_type)
        if os.path.isdir(eclipse_val_dir):
            val_counts[eclipse_type] = len(os.listdir(eclipse_val_dir))

    return val_counts

def test_train_val_split_correct_ratio(implementation, setup_test_env):
    """Test if implementation splits the training data correctly with the specified ratio."""
    impl_name, module = implementation

    # The setup_test_env fixture already creates the directory structure
    data_path = setup_test_env

    # Create val directory structure first to help implementations
    create_val_dir_if_needed(impl_name, module)

    # Create function wrapper for the implementation
    run_function = create_dummy_function(module)

    # Run the implementation
    run_function(data_path, DEFAULT_VAL_FRAC)

    # Check if the validation directory exists
    val_dir = os.path.join(data_path, "val")
    if not os.path.exists(val_dir):
        # Try to create it and run again if needed
        os.makedirs(val_dir, exist_ok=True)
        run_function(data_path, DEFAULT_VAL_FRAC)

    assert os.path.exists(val_dir), "Validation directory not created"

    # Check each eclipse type folder
    train_dir = os.path.join(data_path, "train")
    eclipse_types = [d for d in os.listdir(train_dir) if os.path.isdir(os.path.join(train_dir, d))]

    for eclipse_type in eclipse_types:
        val_eclipse_dir = os.path.join(val_dir, eclipse_type)

        # Create the directory if it doesn't exist
        if not os.path.exists(val_eclipse_dir):
            os.makedirs(val_eclipse_dir, exist_ok=True)
            # Run the implementation again
            run_function(data_path, DEFAULT_VAL_FRAC)

        assert os.path.exists(val_eclipse_dir), f"Validation directory for {eclipse_type} not created"

        # Count images in train and val
        train_imgs = len(os.listdir(os.path.join(train_dir, eclipse_type)))
        val_imgs = len(os.listdir(val_eclipse_dir))
        total_imgs = train_imgs + val_imgs

        # Skip if no validation images were created
        if val_imgs == 0:
            continue

        # Check if the split ratio is close to VAL_FRAC
        # Allow for minor rounding differences
        expected_val_count = int(100 * DEFAULT_VAL_FRAC)  # 10 total images with 30% in validation
        # Some implementations might add one image to validation if the calculation gives 0
        assert val_imgs in [expected_val_count, expected_val_count + 1], \
            f"Expected approximately {expected_val_count} validation images, got {val_imgs}"
        assert train_imgs + val_imgs == 100, f"Expected 10 total images, got {train_imgs + val_imgs}"

def test_data_integrity(implementation, setup_test_env):
    """Test if the data is properly copied to validation and removed from training."""
    impl_name, module = implementation

    # The setup_test_env fixture already creates the directory structure
    data_path = setup_test_env

    # Create val directory structure first to help implementations
    create_val_dir_if_needed(impl_name, module)

    # Create a list of all original images before splitting
    original_images = {}
    train_dir = os.path.join(data_path, "train")
    for eclipse_type in os.listdir(train_dir):
        eclipse_path = os.path.join(train_dir, eclipse_type)
        if os.path.isdir(eclipse_path):
            original_images[eclipse_type] = set(os.listdir(eclipse_path))

    # Run the implementation
    run_function = create_dummy_function(module)
    run_function(data_path, DEFAULT_VAL_FRAC)

    # Check if files were properly moved/copied
    val_dir = os.path.join(data_path, "val")
    if not os.path.exists(val_dir):
        os.makedirs(val_dir, exist_ok=True)
        run_function(data_path, DEFAULT_VAL_FRAC)

    assert os.path.exists(val_dir), "Validation directory was not created"

    for eclipse_type in original_images:
        # Get current lists of files
        val_eclipse_dir = os.path.join(val_dir, eclipse_type)
        if not os.path.exists(val_eclipse_dir):
            os.makedirs(val_eclipse_dir, exist_ok=True)
            run_function(data_path, DEFAULT_VAL_FRAC)

        if not os.path.exists(val_eclipse_dir):
            continue  # Skip if directory wasn't created after retry

        val_images = set(os.listdir(val_eclipse_dir))
        train_images = set(os.listdir(os.path.join(train_dir, eclipse_type)))

        # If no split happened, skip the test
        if len(val_images) == 0:
            continue

        # Make sure there's no overlap (files should be moved, not duplicated)
        assert len(train_images.intersection(val_images)) == 0, "Files appear in both train and validation"

        # Make sure all original files are accounted for
        assert (train_images.union(val_images)) == original_images[eclipse_type], "Some files are missing after split"

        # Verify content integrity for files in validation
        for img in val_images:
            val_img_path = os.path.join(val_dir, eclipse_type, img)
            with open(val_img_path, "r") as f:
                content = f.read()
            assert content == "dummy image content", "File content was corrupted during copying"

def test_deterministic_split(implementation, tmp_path):
    """Test if the implementation produces deterministic splits with fixed random state."""
    impl_name, module = implementation

    # First run
    test_data_path1 = tmp_path / "test_eclipse_data1"
    test_data_path1.mkdir()

    # Create test environment for first run
    train_dir1 = test_data_path1 / "train"
    train_dir1.mkdir()

    # Create eclipse type directories
    eclipse_types = ["solar", "lunar", "partial"]
    for eclipse_type in eclipse_types:
        eclipse_path = train_dir1 / eclipse_type
        eclipse_path.mkdir()

        # Create dummy image files
        for i in range(100):  # 10 images per type
            img_path = eclipse_path / f"img_{i}.jpg"
            img_path.write_text("dummy image content")

    # Create val directory structure first
    val_dir1 = test_data_path1 / "val"
    val_dir1.mkdir()
    for eclipse_type in eclipse_types:
        (val_dir1 / eclipse_type).mkdir()

    val_images_first_run = {}
    val_counts_first = count_val_images_after_split(module, str(test_data_path1), DEFAULT_VAL_FRAC)

    # Get validation image filenames
    if val_dir1.exists():
        for eclipse_type in os.listdir(val_dir1):
            if (val_dir1 / eclipse_type).is_dir():
                val_images_first_run[eclipse_type] = set(os.listdir(val_dir1 / eclipse_type))

    # Second run
    test_data_path2 = tmp_path / "test_eclipse_data2"
    test_data_path2.mkdir()

    # Create test environment for second run
    train_dir2 = test_data_path2 / "train"
    train_dir2.mkdir()

    for eclipse_type in eclipse_types:
        eclipse_path = train_dir2 / eclipse_type
        eclipse_path.mkdir()

        # Create dummy image files
        for i in range(100):  # 10 images per type
            img_path = eclipse_path / f"img_{i}.jpg"
            img_path.write_text("dummy image content")

    # Create val directory structure first
    val_dir2 = test_data_path2 / "val"
    val_dir2.mkdir()
    for eclipse_type in eclipse_types:
        (val_dir2 / eclipse_type).mkdir()

    val_images_second_run = {}
    val_counts_second = count_val_images_after_split(module, str(test_data_path2), DEFAULT_VAL_FRAC)

    # Get validation image filenames
    if val_dir2.exists():
        for eclipse_type in os.listdir(val_dir2):
            if (val_dir2 / eclipse_type).is_dir():
                val_images_second_run[eclipse_type] = set(os.listdir(val_dir2 / eclipse_type))

    # Skip the test if no validation images in either run
    if not val_counts_first or not val_counts_second:
        return

    # Check if both runs produced the same validation counts at least
    assert val_counts_first == val_counts_second, "Number of validation images is not deterministic"

    # Check if both runs produced the same validation sets
    for eclipse_type in val_images_first_run:
        if eclipse_type in val_images_second_run:
            assert val_images_first_run[eclipse_type] == val_images_second_run[eclipse_type], \
                f"Split is not deterministic for {eclipse_type}"

def test_error_handling(implementation, setup_test_env):
    """Test if implementation handles errors gracefully."""
    impl_name, module = implementation

    # The setup_test_env fixture already creates the directory structure
    data_path = setup_test_env

    # Create val directory structure first to help implementations
    create_val_dir_if_needed(impl_name, module)

    # Create an edge case directory structure
    # Add an empty eclipse type directory
    empty_dir = os.path.join(data_path, "train", "empty_eclipse")
    os.makedirs(empty_dir, exist_ok=True)

    try:
        run_function = create_dummy_function(module)
        run_function(data_path, DEFAULT_VAL_FRAC)

        # Should get here without exceptions
        assert True

        # Check if val directory for empty_eclipse exists
        val_empty_dir = os.path.join(data_path, "val", "empty_eclipse")
        # Some implementations might skip empty directories
        if os.path.exists(val_empty_dir):
            assert os.path.isdir(val_empty_dir), "Validation directory for empty eclipse type not created"
    except Exception as e:
        pytest.fail(f"Implementation failed to handle error gracefully: {str(e)}")

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            # r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output


Instruction: create telegram bot (aiogram 3)

import pytest
import inspect
import asyncio
import re
from unittest.mock import patch, AsyncMock, MagicMock
import logging
from typing import Tuple, Any, List, Dict, Optional

def test_imports_aiogram(implementation):
    """Test that the implementation imports required aiogram components"""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)
    except (TypeError, OSError):
        pytest.skip(f"Could not get source code for {impl_name}")

    # Check for essential aiogram imports
    assert any(pattern in source_code for pattern in [
        "from aiogram import",
        "import aiogram"
    ]), f"{impl_name} should import the aiogram library"

def detect_aiogram_version(source_code: str) -> str:
    """Helper function to detect aiogram version from code patterns"""
    # Aiogram 3 specific patterns
    aiogram3_patterns = [
        r"dp\s*=\s*Dispatcher\(\)",  # No parameters in Dispatcher init
        r"from aiogram\.filters import",  # New filter system
        r"@dp\.message\(",  # New message handler decorator syntax
        r"from aiogram\.enums import",  # Using enums
        r"await dp\.start_polling\(bot\)"  # V3 polling method
    ]

    # Aiogram 2 specific patterns
    aiogram2_patterns = [
        r"@dp\.message_handler",  # Old message handler syntax
        r"dp\s*=\s*Dispatcher\(bot\)",  # Bot parameter in Dispatcher init
        r"executor\.start_polling"  # Old polling method
    ]

    is_v3 = any(re.search(pattern, source_code) for pattern in aiogram3_patterns)
    is_v2 = any(re.search(pattern, source_code) for pattern in aiogram2_patterns)

    if is_v3:
        return "v3"
    elif is_v2:
        return "v2"
    else:
        return "unknown"

def test_bot_initialization(implementation):
    """Test that the bot is properly initialized with a token"""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)
    except (TypeError, OSError):
        pytest.skip(f"Could not get source code for {impl_name}")

    # Check for bot initialization with token
    token_patterns = [
        r"Bot\(\s*token=",
        r"Bot\([^,)]*token",
        r"Bot\(['\"][^'\"]+['\"]"  # Some might pass token directly
    ]

    assert any(re.search(pattern, source_code) for pattern in token_patterns), \
        f"{impl_name} should initialize a Bot with a token"

    # Check for token variable definition
    token_var_patterns = [
        r"(?:API_TOKEN|BOT_TOKEN|TOKEN)\s*=",
        r"token\s*="
    ]

    assert any(re.search(pattern, source_code, re.IGNORECASE) for pattern in token_var_patterns), \
        f"{impl_name} should define a token variable (API_TOKEN, BOT_TOKEN, TOKEN, etc.)"

def test_main_polling_setup(implementation):
    """Test that the implementation includes a main function with proper polling setup"""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)
    except (TypeError, OSError):
        pytest.skip(f"Could not get source code for {impl_name}")

    version = detect_aiogram_version(source_code)

    # Check for main function or equivalent entry point
    main_patterns = [
        r"(async\s+)?def\s+main\s*\(",
        r"if\s+__name__\s*==\s*['\"]__main__['\"]",
        r"asyncio\.run\(",
        r"executor\.start_polling"
    ]

    has_main_function = any(re.search(pattern, source_code, re.MULTILINE) for pattern in main_patterns)

    # Check for polling setup based on version
    if version == "v3":
        polling_patterns = [
            r"await dp\.start_polling\(bot",
            r"await dp\.start\s*\(",
            r"dp\.run_polling\("
        ]
    else:  # v2 or unknown
        polling_patterns = [
            r"executor\.start_polling\(dp",
            r"dp\.start_polling\("
        ]

    # Check for asyncio.run pattern for both versions
    asyncio_patterns = [
        r"asyncio\.run\(main\(\)\)",
        r"asyncio\.run\(",
        r"asyncio\.get_event_loop\(\)\.run_until_complete"
    ]

    has_polling = any(re.search(pattern, source_code, re.MULTILINE) for pattern in polling_patterns)
    has_asyncio_run = any(re.search(pattern, source_code, re.MULTILINE) for pattern in asyncio_patterns)

    assert has_main_function, f"{impl_name} should include a main function or entry point"
    assert has_polling or has_asyncio_run, f"{impl_name} should include a proper polling mechanism for the bot"

def test_proper_async_usage(implementation):
    """Test that the implementation properly uses async/await patterns"""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)
    except (TypeError, OSError):
        pytest.skip(f"Could not get source code for {impl_name}")

    # Check for async function definitions
    has_async_def = "async def" in source_code

    # Check for await usage
    has_await = "await" in source_code

    # Check handlers are defined as async
    handler_patterns = [
        r"@dp\.\w+.*\s+async def",  # Generic handler pattern
        r"@dp\.message.*\s+async def",  # v3 message handler
        r"@dp\.message_handler.*\s+async def",  # v2 message handler
        r"async def \w+\s*\(\s*message:"  # Fallback for non-decorated handlers
    ]

    handlers_async = any(re.search(pattern, source_code, re.MULTILINE) for pattern in handler_patterns)

    assert has_async_def, f"{impl_name} should define async functions"
    assert has_await, f"{impl_name} should use await for async calls"
    assert handlers_async, f"{impl_name} should define message handlers as async functions"

def test_error_handling(implementation):
    """Test that the implementation includes error handling or proper finalization"""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)
    except (TypeError, OSError):
        pytest.skip(f"Could not get source code for {impl_name}")

    # Look for error handling patterns
    error_handling_patterns = [
        r"try\s*:",
        r"except\s+",
        r"finally\s*:",
        r"(?:bot|session)\.(?:close|session\.close)\(\)",
        r"logging\.basicConfig",
        r"logging\.(?:info|error|warning|debug|critical)",
        r"register_errors_handler",
        r"@dp\.errors_handler",
        r"@dp\.error",
        r"print\(.*[Ee]rror"  # Simple error printing
    ]

    has_error_handling = any(re.search(pattern, source_code, re.MULTILINE)
                             for pattern in error_handling_patterns)

    assert has_error_handling, \
        f"{impl_name} should include error handling, session cleanup, or logging"

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            # r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output