backend/tests/test_experiment_tracking.py

"""
Comprehensive tests for experiment tracking system.

Tests experiment creation, status updates, result storage, retrieval,
and integration with evaluation service.
"""

import pytest
import uuid
import json
from datetime import datetime, timedelta
from unittest.mock import Mock, patch, AsyncMock
from typing import Dict, Any

from experiments.experiment_manager import ExperimentManager, get_experiment_manager
from experiments.experiment_store import ExperimentStore, SyncExperimentStore
from experiments.experiment_utils import (
    ExperimentAnalyzer, ExperimentExporter, 
    create_config_snapshot_from_pipeline_config,
    create_result_summary_from_pipeline_result
)
from schemas.experiment_schema import (
    Experiment, ExperimentCreate, ExperimentUpdate, ExperimentStatus,
    ExperimentPriority, ConfigSnapshot, ResultSummary, ExperimentFilter,
    ExperimentList, ExperimentStats
)


class TestExperimentStore:
    """Test the experiment store functionality."""
    
    @pytest.fixture
    def mock_store(self):
        """Create a mock experiment store."""
        store = Mock(spec=ExperimentStore)
        store.create_experiment = AsyncMock()
        store.get_experiment = AsyncMock()
        store.update_experiment = AsyncMock()
        store.list_experiments = AsyncMock()
        store.delete_experiment = AsyncMock()
        store.get_experiment_stats = AsyncMock()
        return store
    
    @pytest.fixture
    def sample_experiment(self):
        """Create a sample experiment for testing."""
        config_snapshot = ConfigSnapshot(
            model_name="test-model",
            model_config={"param": "value"},
            dataset_name="jailbreak",
            dataset_version="v1",
            attack_types=["jailbreak", "injection"],
            pipeline_config={"max_iterations": 5},
            prompt_count=10,
            max_iterations=5,
            mutation_enabled=True,
            weights={"robustness": 0.3},
            environment={}
        )
        
        return Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=config_snapshot,
            model_name="test-model",
            dataset_name="jailbreak",
            dataset_version="v1",
            attack_types=["jailbreak", "injection"],
            prompt_count=10,
            status=ExperimentStatus.PENDING
        )
    
    @pytest.mark.asyncio
    async def test_create_experiment(self, mock_store, sample_experiment):
        """Test experiment creation."""
        mock_store.create_experiment.return_value = sample_experiment
        
        result = await mock_store.create_experiment(sample_experiment)
        
        assert result == sample_experiment
        mock_store.create_experiment.assert_called_once_with(sample_experiment)
    
    @pytest.mark.asyncio
    async def test_get_experiment(self, mock_store, sample_experiment):
        """Test experiment retrieval."""
        mock_store.get_experiment.return_value = sample_experiment
        
        result = await mock_store.get_experiment(str(sample_experiment.run_id))
        
        assert result == sample_experiment
        mock_store.get_experiment.assert_called_once_with(str(sample_experiment.run_id))
    
    @pytest.mark.asyncio
    async def test_update_experiment(self, mock_store, sample_experiment):
        """Test experiment update."""
        update_data = ExperimentUpdate(status=ExperimentStatus.RUNNING)
        updated_experiment = sample_experiment.copy()
        updated_experiment.status = ExperimentStatus.RUNNING
        
        mock_store.update_experiment.return_value = updated_experiment
        
        result = await mock_store.update_experiment(str(sample_experiment.run_id), update_data)
        
        assert result.status == ExperimentStatus.RUNNING
        mock_store.update_experiment.assert_called_once_with(str(sample_experiment.run_id), update_data)
    
    @pytest.mark.asyncio
    async def test_list_experiments(self, mock_store):
        """Test experiment listing."""
        filters = ExperimentFilter(page=1, page_size=20)
        experiment_list = ExperimentList(
            experiments=[],
            total_count=0,
            page=1,
            page_size=20,
            total_pages=0
        )
        
        mock_store.list_experiments.return_value = experiment_list
        
        result = await mock_store.list_experiments(filters)
        
        assert result.total_count == 0
        mock_store.list_experiments.assert_called_once_with(filters)
    
    @pytest.mark.asyncio
    async def test_delete_experiment(self, mock_store):
        """Test experiment deletion."""
        mock_store.delete_experiment.return_value = True
        
        result = await mock_store.delete_experiment("test-run-id")
        
        assert result is True
        mock_store.delete_experiment.assert_called_once_with("test-run-id")


class TestExperimentManager:
    """Test the experiment manager functionality."""
    
    @pytest.fixture
    def mock_store(self):
        """Create a mock experiment store."""
        store = Mock(spec=ExperimentStore)
        store.create_experiment = AsyncMock()
        store.get_experiment = AsyncMock()
        store.update_experiment = AsyncMock()
        store.list_experiments = AsyncMock()
        store.delete_experiment = AsyncMock()
        store.get_experiment_stats = AsyncMock()
        return store
    
    @pytest.fixture
    def manager(self, mock_store):
        """Create experiment manager with mock store."""
        return ExperimentManager(mock_store)
    
    @pytest.fixture
    def sample_config_snapshot(self):
        """Create a sample config snapshot."""
        return ConfigSnapshot(
            model_name="test-model",
            model_config={"param": "value"},
            dataset_name="jailbreak",
            dataset_version="v1",
            attack_types=["jailbreak"],
            pipeline_config={"max_iterations": 5},
            prompt_count=10,
            max_iterations=5,
            mutation_enabled=True,
            weights={"robustness": 0.3},
            environment={}
        )
    
    @pytest.fixture
    def sample_experiment_create(self, sample_config_snapshot):
        """Create a sample experiment create request."""
        return ExperimentCreate(
            experiment_name="Test Experiment",
            description="Test description",
            config_snapshot=sample_config_snapshot,
            priority=ExperimentPriority.MEDIUM,
            tags=["test"]
        )
    
    def test_create_experiment_success(self, manager, mock_store, sample_experiment_create):
        """Test successful experiment creation."""
        # Mock the store response
        created_experiment = Experiment(
            run_id=uuid.uuid4(),
            experiment_name="Test Experiment",
            description="Test description",
            config_snapshot=sample_experiment_create.config_snapshot,
            model_name="test-model",
            dataset_name="jailbreak",
            dataset_version="v1",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.PENDING,
            priority=ExperimentPriority.MEDIUM,
            tags=["test"]
        )
        
        mock_store.create_experiment.return_value = created_experiment
        
        result = manager.create_experiment(sample_experiment_create)
        
        assert result.experiment_name == "Test Experiment"
        assert result.status == ExperimentStatus.PENDING
        assert result.run_id is not None
        mock_store.create_experiment.assert_called_once()
    
    def test_update_status_success(self, manager, mock_store):
        """Test successful status update."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.PENDING
        )
        
        # Mock get_experiment to return the experiment
        mock_store.get_experiment.return_value = experiment
        
        # Mock update_experiment to return updated experiment
        updated_experiment = experiment.copy()
        updated_experiment.status = ExperimentStatus.RUNNING
        updated_experiment.started_at = datetime.utcnow()
        mock_store.update_experiment.return_value = updated_experiment
        
        result = manager.update_status(str(experiment.run_id), ExperimentStatus.RUNNING)
        
        assert result.status == ExperimentStatus.RUNNING
        assert result.started_at is not None
        mock_store.get_experiment.assert_called_once()
        mock_store.update_experiment.assert_called_once()
    
    def test_update_status_invalid_transition(self, manager, mock_store):
        """Test invalid status transition."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.COMPLETED  # Terminal state
        )
        
        mock_store.get_experiment.return_value = experiment
        
        with pytest.raises(ValueError, match="Invalid status transition"):
            manager.update_status(str(experiment.run_id), ExperimentStatus.RUNNING)
    
    def test_save_results_success(self, manager, mock_store):
        """Test successful result saving."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.RUNNING
        )
        
        result_summary = ResultSummary(
            robustness_score=0.8,
            risk_score=0.2,
            success_rate=0.3,
            total_attacks=10,
            successful_attacks=3,
            failed_attacks=7,
            execution_time_ms=5000
        )
        
        # Mock get_experiment to return the experiment
        mock_store.get_experiment.return_value = experiment
        
        # Mock update_experiment to return updated experiment
        updated_experiment = experiment.copy()
        updated_experiment.status = ExperimentStatus.COMPLETED
        updated_experiment.result_summary = result_summary
        updated_experiment.completed_at = datetime.utcnow()
        mock_store.update_experiment.return_value = updated_experiment
        
        result = manager.save_results(str(experiment.run_id), result_summary)
        
        assert result.status == ExperimentStatus.COMPLETED
        assert result.result_summary == result_summary
        assert result.completed_at is not None
        mock_store.get_experiment.assert_called_once()
        mock_store.update_experiment.assert_called_once()
    
    def test_get_experiment_success(self, manager, mock_store):
        """Test successful experiment retrieval."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.PENDING
        )
        
        mock_store.get_experiment.return_value = experiment
        
        result = manager.get_experiment(str(experiment.run_id))
        
        assert result == experiment
        mock_store.get_experiment.assert_called_once_with(str(experiment.run_id))
    
    def test_get_experiment_not_found(self, manager, mock_store):
        """Test experiment not found."""
        mock_store.get_experiment.return_value = None
        
        result = manager.get_experiment("non-existent-id")
        
        assert result is None
        mock_store.get_experiment.assert_called_once_with("non-existent-id")
    
    def test_mark_as_failed(self, manager, mock_store):
        """Test marking experiment as failed."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.RUNNING
        )
        
        error_message = "Test error message"
        
        # Mock get_experiment to return the experiment
        mock_store.get_experiment.return_value = experiment
        
        # Mock update_experiment to return updated experiment
        updated_experiment = experiment.copy()
        updated_experiment.status = ExperimentStatus.FAILED
        updated_experiment.error_message = error_message
        updated_experiment.completed_at = datetime.utcnow()
        mock_store.update_experiment.return_value = updated_experiment
        
        result = manager.mark_as_failed(str(experiment.run_id), error_message)
        
        assert result.status == ExperimentStatus.FAILED
        assert result.error_message == error_message
        assert result.completed_at is not None
        mock_store.get_experiment.assert_called_once()
        mock_store.update_experiment.assert_called_once()
    
    def test_clone_experiment(self, manager, mock_store):
        """Test experiment cloning."""
        original_experiment = Experiment(
            run_id=uuid.uuid4(),
            experiment_name="Original Experiment",
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.COMPLETED
        )
        
        # Mock get_experiment to return the original
        mock_store.get_experiment.return_value = original_experiment
        
        # Mock create_experiment to return the clone
        cloned_experiment = Experiment(
            run_id=uuid.uuid4(),
            experiment_name="Original Experiment (Clone)",
            config_snapshot=original_experiment.config_snapshot,
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.PENDING,
            parent_experiment_id=original_experiment.run_id
        )
        mock_store.create_experiment.return_value = cloned_experiment
        
        result = manager.clone_experiment(str(original_experiment.run_id))
        
        assert result.experiment_name == "Original Experiment (Clone)"
        assert result.status == ExperimentStatus.PENDING
        assert result.parent_experiment_id == original_experiment.run_id
        mock_store.get_experiment.assert_called_once()
        mock_store.create_experiment.assert_called_once()


class TestExperimentAnalyzer:
    """Test the experiment analyzer functionality."""
    
    @pytest.fixture
    def analyzer(self):
        """Create experiment analyzer."""
        return ExperimentAnalyzer()
    
    @pytest.fixture
    def mock_manager(self):
        """Create mock experiment manager."""
        manager = Mock(spec=ExperimentManager)
        manager.list_experiments = Mock()
        return manager
    
    def test_analyze_performance_trends_no_experiments(self, analyzer, mock_manager):
        """Test performance trends with no experiments."""
        # Mock empty experiment list
        experiment_list = ExperimentList(
            experiments=[],
            total_count=0,
            page=1,
            page_size=20,
            total_pages=0
        )
        mock_manager.list_experiments.return_value = experiment_list
        
        analyzer.manager = mock_manager
        
        result = analyzer.analyze_performance_trends("test-model", 30)
        
        assert "error" in result
        assert "No completed experiments found" in result["error"]
    
    def test_analyze_performance_trends_with_data(self, analyzer, mock_manager):
        """Test performance trends with experiment data."""
        # Create sample experiments
        experiments = []
        for i in range(5):
            result_summary = ResultSummary(
                robustness_score=0.7 + (i * 0.05),
                risk_score=0.3 - (i * 0.02),
                success_rate=0.4 + (i * 0.1),
                total_attacks=10,
                successful_attacks=4 + i,
                failed_attacks=6 - i,
                execution_time_ms=5000 + (i * 100)
            )
            
            experiment = Mock(spec=Experiment)
            experiment.status = ExperimentStatus.COMPLETED
            experiment.result_summary = result_summary
            experiment.created_at = datetime.utcnow() - timedelta(days=i)
            experiments.append(experiment)
        
        experiment_list = ExperimentList(
            experiments=experiments,
            total_count=5,
            page=1,
            page_size=20,
            total_pages=1
        )
        mock_manager.list_experiments.return_value = experiment_list
        
        analyzer.manager = mock_manager
        
        result = analyzer.analyze_performance_trends("test-model", 30)
        
        assert result["model_name"] == "test-model"
        assert result["total_experiments"] == 5
        assert "trends" in result
        assert "current_performance" in result
        assert "performance_history" in result
        assert len(result["performance_history"]) == 5
    
    def test_compare_model_performance(self, analyzer, mock_manager):
        """Test model performance comparison."""
        # Mock experiment lists for different models
        def mock_list_experiments(filters):
            if filters.model_name == "model-a":
                experiment_list = ExperimentList(
                    experiments=[Mock(result_summary=Mock(robustness_score=0.8, risk_score=0.2, success_rate=0.3))],
                    total_count=1,
                    page=1,
                    page_size=20,
                    total_pages=1
                )
            elif filters.model_name == "model-b":
                experiment_list = ExperimentList(
                    experiments=[Mock(result_summary=Mock(robustness_score=0.7, risk_score=0.3, success_rate=0.4))],
                    total_count=1,
                    page=1,
                    page_size=20,
                    total_pages=1
                )
            else:
                experiment_list = ExperimentList(
                    experiments=[],
                    total_count=0,
                    page=1,
                    page_size=20,
                    total_pages=0
                )
            return experiment_list
        
        mock_manager.list_experiments.side_effect = mock_list_experiments
        analyzer.manager = mock_manager
        
        result = analyzer.compare_model_performance(["model-a", "model-b"])
        
        assert result["models_compared"] == ["model-a", "model-b"]
        assert "model_stats" in result
        assert "best_performing_model" in result
        assert "model-a" in result["model_stats"]
        assert "model-b" in result["model_stats"]
    
    def test_analyze_attack_effectiveness(self, analyzer, mock_manager):
        """Test attack effectiveness analysis."""
        # Mock experiments for different attack types
        def mock_list_experiments(filters):
            if filters.attack_types == ["jailbreak"]:
                experiment_list = ExperimentList(
                    experiments=[Mock(result_summary=Mock(success_rate=0.3, total_attacks=10, successful_attacks=3))],
                    total_count=1,
                    page=1,
                    page_size=20,
                    total_pages=1
                )
            elif filters.attack_types == ["injection"]:
                experiment_list = ExperimentList(
                    experiments=[Mock(result_summary=Mock(success_rate=0.2, total_attacks=10, successful_attacks=2))],
                    total_count=1,
                    page=1,
                    page_size=20,
                    total_pages=1
                )
            else:
                experiment_list = ExperimentList(
                    experiments=[],
                    total_count=0,
                    page=1,
                    page_size=20,
                    total_pages=0
                )
            return experiment_list
        
        mock_manager.list_experiments.side_effect = mock_list_experiments
        analyzer.manager = mock_manager
        
        result = analyzer.analyze_attack_effectiveness(["jailbreak", "injection"])
        
        assert result["attack_types_analyzed"] == ["jailbreak", "injection"]
        assert "attack_stats" in result
        assert "most_effective_attack" in result
        assert "least_effective_attack" in result
        assert "jailbreak" in result["attack_stats"]
        assert "injection" in result["attack_stats"]


class TestExperimentExporter:
    """Test the experiment exporter functionality."""
    
    @pytest.fixture
    def exporter(self):
        """Create experiment exporter."""
        return ExperimentExporter()
    
    @pytest.fixture
    def mock_manager(self):
        """Create mock experiment manager."""
        manager = Mock(spec=ExperimentManager)
        manager.get_experiment = Mock()
        manager.list_experiments = Mock()
        return manager
    
    def test_export_experiment_to_dict(self, exporter, mock_manager):
        """Test exporting experiment to dictionary."""
        experiment = Experiment(
            run_id=uuid.uuid4(),
            experiment_name="Test Experiment",
            config_snapshot=Mock(),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.COMPLETED
        )
        
        mock_manager.get_experiment.return_value = experiment
        exporter.manager = mock_manager
        
        result = exporter.export_experiment_to_dict(str(experiment.run_id))
        
        assert result is not None
        assert result["run_id"] == str(experiment.run_id)
        assert result["experiment_name"] == "Test Experiment"
        assert result["model_name"] == "test-model"
    
    def test_export_experiment_to_dict_not_found(self, exporter, mock_manager):
        """Test exporting non-existent experiment."""
        mock_manager.get_experiment.return_value = None
        exporter.manager = mock_manager
        
        result = exporter.export_experiment_to_dict("non-existent-id")
        
        assert result is None
    
    def test_export_experiments_to_csv(self, exporter, mock_manager):
        """Test exporting experiments to CSV."""
        experiments = [
            Mock(
                run_id=uuid.uuid4(),
                experiment_name="Test 1",
                model_name="model-a",
                dataset_name="jailbreak",
                dataset_version="v1",
                status=ExperimentStatus.COMPLETED,
                priority=ExperimentPriority.MEDIUM,
                created_at=datetime.utcnow(),
                completed_at=datetime.utcnow(),
                result_summary=Mock(
                    robustness_score=0.8,
                    risk_score=0.2,
                    success_rate=0.3,
                    total_attacks=10,
                    execution_time_ms=5000
                )
            )
        ]
        
        experiment_list = ExperimentList(
            experiments=experiments,
            total_count=1,
            page=1,
            page_size=20,
            total_pages=1
        )
        
        mock_manager.list_experiments.return_value = experiment_list
        exporter.manager = mock_manager
        
        result = exporter.export_experiments_to_csv()
        
        assert result is not None
        assert "run_id" in result
        assert "experiment_name" in result
        assert "model_name" in result
        assert len(result.split('\n')) > 1  # Header + data
    
    def test_export_experiment_summary(self, exporter, mock_manager):
        """Test exporting experiment summary."""
        result_summary = ResultSummary(
            robustness_score=0.8,
            risk_score=0.2,
            success_rate=0.3,
            total_attacks=10,
            successful_attacks=3,
            failed_attacks=7,
            execution_time_ms=5000
        )
        
        experiment = Experiment(
            run_id=uuid.uuid4(),
            experiment_name="Test Experiment",
            description="Test description",
            config_snapshot=Mock(
                model_name="test-model",
                dataset_name="jailbreak",
                dataset_version="v1",
                attack_types=["jailbreak"],
                prompt_count=10,
                max_iterations=5,
                mutation_enabled=True
            ),
            model_name="test-model",
            attack_types=["jailbreak"],
            prompt_count=10,
            status=ExperimentStatus.COMPLETED,
            priority=ExperimentPriority.MEDIUM,
            created_at=datetime.utcnow(),
            started_at=datetime.utcnow(),
            completed_at=datetime.utcnow(),
            result_summary=result_summary,
            tags=["test"],
            created_by="test@example.com"
        )
        
        mock_manager.get_experiment.return_value = experiment
        exporter.manager = mock_manager
        
        result = exporter.export_experiment_summary(str(experiment.run_id))
        
        assert result is not None
        assert "experiment_info" in result
        assert "configuration" in result
        assert "results" in result
        assert result["experiment_info"]["run_id"] == str(experiment.run_id)
        assert result["configuration"]["model_name"] == "test-model"
        assert result["results"]["robustness_score"] == 0.8


class TestUtilityFunctions:
    """Test utility functions."""
    
    def test_create_config_snapshot_from_pipeline_config(self):
        """Test creating config snapshot from pipeline config."""
        pipeline_config = Mock()
        pipeline_config.dict.return_value = {
            "max_iterations": 5,
            "attack_types": ["jailbreak"],
            "num_prompts": 10,
            "mutation_enabled": True,
            "weights": {"robustness": 0.3}
        }
        pipeline_config.dataset_name = "jailbreak"
        pipeline_config.dataset_version = "v1"
        pipeline_config.attack_types = ["jailbreak"]
        pipeline_config.num_prompts = 10
        pipeline_config.max_iterations = 5
        pipeline_config.mutation_enabled = True
        pipeline_config.weights = {"robustness": 0.3}
        
        result = create_config_snapshot_from_pipeline_config(
            pipeline_config, 
            "test-model",
            {"param": "value"}
        )
        
        assert isinstance(result, ConfigSnapshot)
        assert result.model_name == "test-model"
        assert result.dataset_name == "jailbreak"
        assert result.attack_types == ["jailbreak"]
        assert result.prompt_count == 10
        assert result.max_iterations == 5
        assert result.mutation_enabled is True
    
    def test_create_result_summary_from_pipeline_result(self):
        """Test creating result summary from pipeline result."""
        score_card = Mock()
        score_card.robustness_score = 0.8
        score_card.risk_score = 0.2
        score_card.hallucination_rate = 0.1
        score_card.toxicity_rate = 0.05
        score_card.confidence_score = 0.9
        
        pipeline_result = Mock()
        pipeline_result.score_card = score_card
        pipeline_result.success_rate = 0.3
        pipeline_result.total_attacks = 10
        pipeline_result.successful_attacks = [Mock(), Mock(), Mock()]
        pipeline_result.failed_attacks = [Mock(), Mock(), Mock(), Mock(), Mock(), Mock(), Mock()]
        pipeline_result.execution_time_ms = 5000
        
        result = create_result_summary_from_pipeline_result(pipeline_result)
        
        assert isinstance(result, ResultSummary)
        assert result.robustness_score == 0.8
        assert result.risk_score == 0.2
        assert result.success_rate == 0.3
        assert result.total_attacks == 10
        assert result.successful_attacks == 3
        assert result.failed_attacks == 7
        assert result.execution_time_ms == 5000


class TestGlobalFunctions:
    """Test global convenience functions."""
    
    @patch('experiments.experiment_manager._experiment_manager')
    def test_get_experiment_manager(self, mock_global_manager):
        """Test global experiment manager getter."""
        from experiments.experiment_manager import get_experiment_manager
        
        manager_instance = Mock()
        mock_global_manager.return_value = manager_instance
        
        result = get_experiment_manager()
        
        assert result == manager_instance
        mock_global_manager.assert_called_once()
    
    @patch('experiments.experiment_manager.get_experiment_manager')
    def test_create_experiment(self, mock_get_manager):
        """Test global create_experiment function."""
        from experiments.experiment_manager import create_experiment
        
        manager = Mock()
        mock_get_manager.return_value = manager
        
        config = Mock(spec=ExperimentCreate)
        experiment = Mock(spec=Experiment)
        manager.create_experiment.return_value = experiment
        
        result = create_experiment(config)
        
        assert result == experiment
        manager.create_experiment.assert_called_once_with(config)
    
    @patch('experiments.experiment_manager.get_experiment_manager')
    def test_update_experiment_status(self, mock_get_manager):
        """Test global update_experiment_status function."""
        from experiments.experiment_manager import update_experiment_status
        
        manager = Mock()
        mock_get_manager.return_value = manager
        
        experiment = Mock(spec=Experiment)
        manager.update_status.return_value = experiment
        
        result = update_experiment_status("test-id", ExperimentStatus.RUNNING)
        
        assert result == experiment
        manager.update_status.assert_called_once_with("test-id", ExperimentStatus.RUNNING)
    
    @patch('experiments.experiment_manager.get_experiment_manager')
    def test_save_experiment_results(self, mock_get_manager):
        """Test global save_experiment_results function."""
        from experiments.experiment_manager import save_experiment_results
        
        manager = Mock()
        mock_get_manager.return_value = manager
        
        experiment = Mock(spec=Experiment)
        result_summary = Mock(spec=ResultSummary)
        full_result = {"test": "data"}
        manager.save_results.return_value = experiment
        
        result = save_experiment_results("test-id", result_summary, full_result)
        
        assert result == experiment
        manager.save_results.assert_called_once_with("test-id", result_summary, full_result)
    
    @patch('experiments.experiment_manager.get_experiment_manager')
    def test_get_experiment_by_id(self, mock_get_manager):
        """Test global get_experiment_by_id function."""
        from experiments.experiment_manager import get_experiment_by_id
        
        manager = Mock()
        mock_get_manager.return_value = manager
        
        experiment = Mock(spec=Experiment)
        manager.get_experiment.return_value = experiment
        
        result = get_experiment_by_id("test-id")
        
        assert result == experiment
        manager.get_experiment.assert_called_once_with("test-id")


if __name__ == "__main__":
    pytest.main([__file__])