""" Pytest fixtures for runpod tests. Provides mocked clients, test data, and reusable test setup. """ import pytest from unittest.mock import Mock, MagicMock from pathlib import Path import tempfile import json @pytest.fixture def mock_env(monkeypatch): """Set up mock environment variables.""" env_vars = { "RUNPOD_API_KEY": "test_api_key_" + "x" * 32, "RUNPOD_VOLUME_ID": "se3zdnb5o4", "RUNPOD_CONTAINER_REGISTRY_AUTH_ID": "test_registry_auth" } for key, value in env_vars.items(): monkeypatch.setenv(key, value) return env_vars @pytest.fixture def sample_config(): """Sample RunPodConfig for testing.""" from runpod.config import RunPodConfig return RunPodConfig( api_key="test_api_key_" + "x" * 32, volume_id="se3zdnb5o4", container_registry_auth_id="test_registry_auth", datacenters=["EUR-IS-1"] ) @pytest.fixture def mock_gpu_data(): """Sample GPU data from RunPod API.""" return { "data": { "gpuTypes": [ { "id": "NVIDIA RTX A4000", "displayName": "RTX A4000", "memoryInGb": 16, "secureCloud": 10, "communityCloud": 0, "lowestPrice": { "uninterruptablePrice": "0.25" } }, { "id": "NVIDIA Tesla V100", "displayName": "Tesla V100", "memoryInGb": 16, "secureCloud": 5, "communityCloud": 0, "lowestPrice": { "uninterruptablePrice": "0.10" } }, { "id": "NVIDIA RTX 3090", "displayName": "RTX 3090", "memoryInGb": 24, "secureCloud": 0, # Not available in secure cloud "communityCloud": 20, "lowestPrice": { "uninterruptablePrice": "0.15" } }, { "id": "NVIDIA Tesla T4", "displayName": "Tesla T4", "memoryInGb": 15, # Below 16GB threshold "secureCloud": 8, "communityCloud": 0, "lowestPrice": { "uninterruptablePrice": "0.08" } } ] } } @pytest.fixture def mock_pod_data(): """Sample pod deployment response.""" return { "id": "test-pod-123456", "desiredStatus": "RUNNING", "imageName": "jgrusewski/foxhunt:latest", "containerDiskInGb": 50, "costPerHr": "0.25", "machine": { "gpuType": { "id": "NVIDIA RTX A4000", "displayName": "RTX A4000" }, "dataCenterId": "EUR-IS-1" }, "gpu": { "count": 1 } } @pytest.fixture def mock_requests_session(): """Mock requests session for API calls.""" session = MagicMock() session.headers = {} return session @pytest.fixture def mock_s3_client(): """Mock boto3 S3 client.""" client = MagicMock() client.upload_file = MagicMock(return_value=None) client.download_file = MagicMock(return_value=None) client.get_object = MagicMock() client.list_objects_v2 = MagicMock() client.delete_object = MagicMock(return_value=None) return client @pytest.fixture def temp_directory(): """Create temporary directory for file operations.""" with tempfile.TemporaryDirectory() as tmpdir: yield Path(tmpdir) @pytest.fixture def sample_binary_files(temp_directory): """Create sample binary files for upload testing.""" binaries_dir = temp_directory / "binaries" binaries_dir.mkdir() # Create sample binaries binary_names = [ "train_tft_parquet", "train_mamba2_parquet", "train_dqn", "train_ppo" ] created_files = [] for name in binary_names: binary_path = binaries_dir / name binary_path.write_bytes(b"fake binary data " * 100) created_files.append(binary_path) return binaries_dir, created_files @pytest.fixture def sample_parquet_files(temp_directory): """Create sample parquet test data files.""" data_dir = temp_directory / "test_data" data_dir.mkdir() parquet_files = [ "ES_FUT_180d.parquet", "NQ_FUT_180d.parquet" ] created_files = [] for name in parquet_files: parquet_path = data_dir / name parquet_path.write_bytes(b"fake parquet data " * 100) created_files.append(parquet_path) return data_dir, created_files @pytest.fixture def mock_log_content(): """Sample log content for log tailing tests.""" return """ [2025-10-29 12:00:00] Starting training... [2025-10-29 12:00:01] Loading data from /runpod-volume/test_data/ES_FUT_180d.parquet [2025-10-29 12:00:02] Data loaded: 10000 samples [2025-10-29 12:00:03] Epoch 1/50 - Loss: 0.5234 [2025-10-29 12:01:00] Epoch 2/50 - Loss: 0.4123 [2025-10-29 12:02:00] Epoch 3/50 - Loss: 0.3567 [2025-10-29 12:30:00] Training complete [2025-10-29 12:30:01] Model saved successfully to /runpod-volume/models/tft_best.safetensors [2025-10-29 12:30:02] Checkpoint saved """ @pytest.fixture def mock_api_responses(): """Collection of mock API responses for various scenarios.""" return { "success_deployment": { "id": "test-pod-123456", "desiredStatus": "RUNNING", "costPerHr": "0.25", "containerDiskInGb": 50 }, "error_no_availability": { "error": "No machines available in specified datacenters" }, "error_invalid_gpu": { "error": "Invalid GPU type specified" }, "graphql_error": { "errors": [ { "message": "Invalid API key", "extensions": {"code": "UNAUTHENTICATED"} } ] }, "pod_status_running": { "id": "test-pod-123456", "desiredStatus": "RUNNING" }, "pod_status_completed": { "id": "test-pod-123456", "desiredStatus": "COMPLETED" }, "pod_status_failed": { "id": "test-pod-123456", "desiredStatus": "FAILED" } } @pytest.fixture def mock_s3_file_list(): """Mock S3 file listing response.""" return { "Contents": [ { "Key": "models/tft_best.safetensors", "Size": 1024000, "LastModified": "2025-10-29T12:30:00Z" }, { "Key": "models/tft_config.json", "Size": 1024, "LastModified": "2025-10-29T12:30:01Z" }, { "Key": "logs/training.log", "Size": 5120, "LastModified": "2025-10-29T12:30:02Z" } ] } @pytest.fixture def env_file_content(): """Sample .env.runpod file content.""" return """ RUNPOD_API_KEY=test_api_key_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx RUNPOD_VOLUME_ID=se3zdnb5o4 RUNPOD_CONTAINER_REGISTRY_AUTH_ID=test_registry_auth AWS_ACCESS_KEY_ID=test_access_key AWS_SECRET_ACCESS_KEY=test_secret_key RUNPOD_S3_BUCKET=se3zdnb5o4 RUNPOD_S3_ENDPOINT=https://s3api-eur-is-1.runpod.io """.strip() @pytest.fixture def create_env_file(temp_directory, env_file_content): """Create a temporary .env.runpod file.""" env_file = temp_directory / ".env.runpod" env_file.write_text(env_file_content) return env_file