282 lines
7.7 KiB
Python
282 lines
7.7 KiB
Python
"""
|
|
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
|