390 lines
13 KiB
Python
390 lines
13 KiB
Python
"""
|
|
Tests for runpod.client module.
|
|
|
|
Tests GPU query parsing, pod deployment, error handling, and retry logic.
|
|
"""
|
|
|
|
import pytest
|
|
import responses
|
|
from unittest.mock import MagicMock, patch, Mock
|
|
import requests
|
|
|
|
from runpod.client import RunPodClient, RunPodAPIError
|
|
from runpod.config import RunPodConfig
|
|
|
|
|
|
class TestRunPodClient:
|
|
"""Tests for RunPodClient class."""
|
|
|
|
def test_client_initialization(self, sample_config):
|
|
"""Test client initializes with config."""
|
|
client = RunPodClient(sample_config)
|
|
|
|
assert client.config == sample_config
|
|
assert client.session is not None
|
|
assert "Authorization" in client.session.headers
|
|
assert client.session.headers["Authorization"] == f"Bearer {sample_config.api_key}"
|
|
|
|
@responses.activate
|
|
def test_query_graphql_success(self, sample_config, mock_gpu_data):
|
|
"""Test successful GraphQL query."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=mock_gpu_data,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
result = client.query_graphql(RunPodClient.GPU_QUERY)
|
|
|
|
assert result == mock_gpu_data
|
|
assert len(responses.calls) == 1
|
|
|
|
@responses.activate
|
|
def test_query_graphql_with_variables(self, sample_config):
|
|
"""Test GraphQL query with variables."""
|
|
expected_response = {"data": {"test": "value"}}
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=expected_response,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
variables = {"gpuType": "RTX 4090"}
|
|
result = client.query_graphql("query Test($gpuType: String)", variables)
|
|
|
|
assert result == expected_response
|
|
assert len(responses.calls) == 1
|
|
|
|
# Verify variables were sent
|
|
request_body = responses.calls[0].request.body
|
|
assert b'"variables"' in request_body
|
|
|
|
@responses.activate
|
|
def test_query_graphql_error_response(self, sample_config, mock_api_responses):
|
|
"""Test GraphQL query with error response."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=mock_api_responses["graphql_error"],
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="Invalid API key"):
|
|
client.query_graphql(RunPodClient.GPU_QUERY)
|
|
|
|
@responses.activate
|
|
def test_query_graphql_network_error(self, sample_config):
|
|
"""Test GraphQL query with network error."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
body=requests.exceptions.ConnectionError("Network error")
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="Failed to query"):
|
|
client.query_graphql(RunPodClient.GPU_QUERY)
|
|
|
|
@responses.activate
|
|
def test_query_graphql_timeout(self, sample_config):
|
|
"""Test GraphQL query timeout."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
body=requests.exceptions.Timeout("Request timeout")
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError):
|
|
client.query_graphql(RunPodClient.GPU_QUERY, timeout=1)
|
|
|
|
@responses.activate
|
|
def test_get_available_gpus_success(self, sample_config, mock_gpu_data):
|
|
"""Test getting available GPUs."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=mock_gpu_data,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
gpus = client.get_available_gpus(min_vram_gb=16)
|
|
|
|
# Should return 2 GPUs: RTX A4000 and Tesla V100 (secure cloud only)
|
|
assert len(gpus) == 2
|
|
|
|
# Should be sorted by price (V100 cheaper than A4000)
|
|
assert gpus[0]["name"] == "Tesla V100"
|
|
assert gpus[0]["price"] == 0.10
|
|
assert gpus[1]["name"] == "RTX A4000"
|
|
assert gpus[1]["price"] == 0.25
|
|
|
|
@responses.activate
|
|
def test_get_available_gpus_filters_vram(self, sample_config, mock_gpu_data):
|
|
"""Test GPU filtering by VRAM."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=mock_gpu_data,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
gpus = client.get_available_gpus(min_vram_gb=24)
|
|
|
|
# Only RTX 3090 has 24GB, but it's not in secure cloud
|
|
assert len(gpus) == 0
|
|
|
|
@responses.activate
|
|
def test_get_available_gpus_filters_secure_cloud(self, sample_config, mock_gpu_data):
|
|
"""Test GPU filtering by secure cloud availability."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json=mock_gpu_data,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
gpus = client.get_available_gpus(min_vram_gb=16, secure_cloud_only=False)
|
|
|
|
# Should include RTX 3090 (community cloud only)
|
|
assert len(gpus) == 3 # A4000, V100, 3090
|
|
|
|
@responses.activate
|
|
def test_get_available_gpus_no_results(self, sample_config):
|
|
"""Test getting GPUs when none available."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.GRAPHQL_ENDPOINT,
|
|
json={"data": {"gpuTypes": []}},
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
gpus = client.get_available_gpus()
|
|
|
|
assert gpus == []
|
|
|
|
@responses.activate
|
|
def test_deploy_pod_success(self, sample_config, mock_pod_data):
|
|
"""Test successful pod deployment."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
json=mock_pod_data,
|
|
status=201
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
result = client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="jgrusewski/foxhunt:latest",
|
|
command="--epochs 50",
|
|
container_disk_gb=50
|
|
)
|
|
|
|
assert result["id"] == "test-pod-123456"
|
|
assert result["desiredStatus"] == "RUNNING"
|
|
assert len(responses.calls) == 1
|
|
|
|
@responses.activate
|
|
def test_deploy_pod_with_custom_options(self, sample_config, mock_pod_data):
|
|
"""Test pod deployment with custom options."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
json=mock_pod_data,
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
result = client.deploy_pod(
|
|
gpu_id="NVIDIA Tesla V100",
|
|
image="custom/image:v2",
|
|
command="python train.py --config custom",
|
|
container_disk_gb=100,
|
|
datacenters=["EUR-IS-1", "EUR-IS-2"],
|
|
pod_name="custom-pod",
|
|
ports=["8888/http", "6006/http", "22/tcp"]
|
|
)
|
|
|
|
assert result["id"] == "test-pod-123456"
|
|
|
|
# Verify request payload
|
|
request_body = responses.calls[0].request.body
|
|
assert b'"name": "custom-pod"' in request_body
|
|
assert b'"containerDiskInGb": 100' in request_body
|
|
|
|
@responses.activate
|
|
def test_deploy_pod_error_no_availability(self, sample_config, mock_api_responses):
|
|
"""Test pod deployment when GPU not available."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
json=mock_api_responses["error_no_availability"],
|
|
status=400
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="No machines available"):
|
|
client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="test/image:latest"
|
|
)
|
|
|
|
@responses.activate
|
|
def test_deploy_pod_error_invalid_response(self, sample_config):
|
|
"""Test pod deployment with invalid response format."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
json={"invalid": "response"}, # Missing 'id' field
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="missing 'id' field"):
|
|
client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="test/image:latest"
|
|
)
|
|
|
|
@responses.activate
|
|
def test_deploy_pod_network_error(self, sample_config):
|
|
"""Test pod deployment with network error."""
|
|
responses.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
body=requests.exceptions.ConnectionError("Network error")
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="request failed"):
|
|
client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="test/image:latest"
|
|
)
|
|
|
|
@responses.activate
|
|
def test_get_pod_status_success(self, sample_config, mock_api_responses):
|
|
"""Test getting pod status."""
|
|
pod_id = "test-pod-123456"
|
|
responses.add(
|
|
responses.GET,
|
|
f"{RunPodClient.REST_API_URL}/{pod_id}",
|
|
json=mock_api_responses["pod_status_running"],
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
status = client.get_pod_status(pod_id)
|
|
|
|
assert status["id"] == pod_id
|
|
assert status["desiredStatus"] == "RUNNING"
|
|
|
|
@responses.activate
|
|
def test_get_pod_status_error(self, sample_config):
|
|
"""Test getting pod status with error."""
|
|
pod_id = "nonexistent-pod"
|
|
responses.add(
|
|
responses.GET,
|
|
f"{RunPodClient.REST_API_URL}/{pod_id}",
|
|
json={"error": "Pod not found"},
|
|
status=404
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError):
|
|
client.get_pod_status(pod_id)
|
|
|
|
@responses.activate
|
|
def test_stop_pod_success(self, sample_config):
|
|
"""Test stopping a pod."""
|
|
pod_id = "test-pod-123456"
|
|
responses.add(
|
|
responses.DELETE,
|
|
f"{RunPodClient.REST_API_URL}/{pod_id}",
|
|
json={"success": True},
|
|
status=200
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
result = client.stop_pod(pod_id)
|
|
|
|
assert result is True
|
|
|
|
@responses.activate
|
|
def test_stop_pod_error(self, sample_config):
|
|
"""Test stopping pod with error."""
|
|
pod_id = "test-pod-123456"
|
|
responses.add(
|
|
responses.DELETE,
|
|
f"{RunPodClient.REST_API_URL}/{pod_id}",
|
|
json={"error": "Failed to stop pod"},
|
|
status=500
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
with pytest.raises(RunPodAPIError, match="Failed to stop"):
|
|
client.stop_pod(pod_id)
|
|
|
|
def test_deploy_pod_command_parsing(self, sample_config, mock_pod_data):
|
|
"""Test that command is properly split into arguments."""
|
|
with patch('runpod.client.requests.Session') as mock_session_class:
|
|
mock_session = MagicMock()
|
|
mock_session_class.return_value = mock_session
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 201
|
|
mock_response.json.return_value = mock_pod_data
|
|
mock_session.post.return_value = mock_response
|
|
|
|
client = RunPodClient(sample_config)
|
|
|
|
client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="test/image:latest",
|
|
command='--parquet-file "/path/with spaces/file.parquet" --epochs 50'
|
|
)
|
|
|
|
# Verify the command was split correctly
|
|
call_args = mock_session.post.call_args
|
|
payload = call_args[1]['json']
|
|
assert 'dockerStartCmd' in payload
|
|
# Should be split into 4 arguments
|
|
assert len(payload['dockerStartCmd']) == 4
|
|
|
|
@pytest.mark.parametrize("status_code", [200, 201])
|
|
def test_deploy_pod_accepts_both_success_codes(self, sample_config, mock_pod_data, status_code):
|
|
"""Test that both 200 and 201 are treated as success."""
|
|
with responses.RequestsMock() as rsps:
|
|
rsps.add(
|
|
responses.POST,
|
|
RunPodClient.REST_API_URL,
|
|
json=mock_pod_data,
|
|
status=status_code
|
|
)
|
|
|
|
client = RunPodClient(sample_config)
|
|
result = client.deploy_pod(
|
|
gpu_id="NVIDIA RTX A4000",
|
|
image="test/image:latest"
|
|
)
|
|
|
|
assert result["id"] == "test-pod-123456"
|