| """Common test fixtures and configuration."""
|
|
|
| import pytest
|
| import asyncio
|
| import os
|
| from unittest.mock import Mock, MagicMock, patch
|
| from typing import Dict, Any, Generator
|
|
|
|
|
| TEST_ENV_VARS = {
|
| "TAVILY_API_KEY": "tvly-test-key-12345",
|
| "NEBIUS_API_KEY": "test-nebius-key",
|
| "OPENAI_API_KEY": "test-openai-key",
|
| "ANTHROPIC_API_KEY": "test-anthropic-key",
|
| "HUGGINGFACE_API_KEY": "test-hf-key",
|
| "LLM_PROVIDER": "nebius"
|
| }
|
|
|
|
|
| for key, value in TEST_ENV_VARS.items():
|
| os.environ[key] = value
|
|
|
| @pytest.fixture
|
| def mock_tavily_client():
|
| """Mock Tavily client for web search tests."""
|
| mock_client = Mock()
|
| mock_client.search.return_value = {
|
| "results": [
|
| {
|
| "title": "Test Result 1",
|
| "url": "https://example.com/1",
|
| "content": "Test content 1",
|
| "score": 0.9
|
| },
|
| {
|
| "title": "Test Result 2",
|
| "url": "https://example.com/2",
|
| "content": "Test content 2",
|
| "score": 0.8
|
| }
|
| ],
|
| "answer": "Test search summary"
|
| }
|
| return mock_client
|
|
|
| @pytest.fixture
|
| def mock_llm_response():
|
| """Mock LLM completion response."""
|
| return '{"sub_questions": ["Question 1?", "Question 2?", "Question 3?"]}'
|
|
|
| @pytest.fixture
|
| def mock_modal_sandbox():
|
| """Mock Modal sandbox for code execution tests."""
|
| mock_sandbox = Mock()
|
| mock_sandbox.exec.return_value = Mock(stdout="Test output", stderr="", returncode=0)
|
| return mock_sandbox
|
|
|
| @pytest.fixture
|
| def sample_user_request():
|
| """Sample user request for testing."""
|
| return "Create a Python script to analyze CSV data and generate charts"
|
|
|
| @pytest.fixture
|
| def sample_search_results():
|
| """Sample search results for testing."""
|
| return [
|
| {
|
| "title": "Python Data Analysis Tutorial",
|
| "url": "https://example.com/pandas-tutorial",
|
| "content": "Learn how to analyze CSV data with pandas and matplotlib...",
|
| "score": 0.95
|
| },
|
| {
|
| "title": "Chart Generation with Python",
|
| "url": "https://example.com/charts",
|
| "content": "Create stunning charts and visualizations...",
|
| "score": 0.87
|
| }
|
| ]
|
|
|
| @pytest.fixture
|
| def sample_code():
|
| """Sample Python code for testing."""
|
| return '''
|
| import pandas as pd
|
| import matplotlib.pyplot as plt
|
|
|
| # Load data
|
| df = pd.read_csv('data.csv')
|
|
|
| # Generate chart
|
| df.plot(kind='bar')
|
| plt.show()
|
| '''
|
|
|
| @pytest.fixture
|
| def mock_config():
|
| """Mock configuration objects."""
|
| api_config = Mock()
|
| api_config.tavily_api_key = "tvly-test-key"
|
| api_config.llm_provider = "nebius"
|
| api_config.nebius_api_key = "test-nebius-key"
|
|
|
| model_config = Mock()
|
| model_config.get_model_for_provider.return_value = "meta-llama/llama-3.1-8b-instruct"
|
|
|
| return api_config, model_config
|
|
|
| @pytest.fixture
|
| def event_loop():
|
| """Create an event loop for async tests."""
|
| loop = asyncio.new_event_loop()
|
| yield loop
|
| loop.close()
|
|
|
| class MockAgent:
|
| """Base mock agent class for testing."""
|
| def __init__(self, name: str):
|
| self.name = name
|
| self.call_count = 0
|
|
|
| def __call__(self, *args, **kwargs):
|
| self.call_count += 1
|
| return {"success": True, "agent": self.name, "calls": self.call_count}
|
|
|
| @pytest.fixture
|
| def mock_agents():
|
| """Mock agent instances for orchestrator testing."""
|
| return {
|
| "question_enhancer": MockAgent("question_enhancer"),
|
| "web_search": MockAgent("web_search"),
|
| "llm_processor": MockAgent("llm_processor"),
|
| "citation_formatter": MockAgent("citation_formatter"),
|
| "code_generator": MockAgent("code_generator"),
|
| "code_runner": MockAgent("code_runner")
|
| }
|
|
|
| @pytest.fixture
|
| def disable_advanced_features():
|
| """Disable advanced features for basic testing."""
|
| with patch('app.ADVANCED_FEATURES_AVAILABLE', False):
|
| yield |
