""" 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__])