440 lines
17 KiB
Python
440 lines
17 KiB
Python
"""
|
|
Tests for runpod.s3_client module.
|
|
|
|
Tests binary upload, log streaming, and result download.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch, Mock
|
|
from pathlib import Path
|
|
from botocore.exceptions import ClientError, BotoCoreError
|
|
import io
|
|
|
|
from runpod.s3_client import S3Client, S3Error
|
|
|
|
|
|
class TestS3Client:
|
|
"""Tests for S3Client class."""
|
|
|
|
@pytest.fixture
|
|
def s3_config(self):
|
|
"""S3 configuration for testing."""
|
|
return {
|
|
"bucket_name": "se3zdnb5o4",
|
|
"endpoint_url": "https://s3api-eur-is-1.runpod.io",
|
|
"aws_access_key_id": "test_access_key",
|
|
"aws_secret_access_key": "test_secret_key",
|
|
"region_name": "eur-is-1"
|
|
}
|
|
|
|
def test_s3_client_initialization(self, s3_config):
|
|
"""Test S3 client initializes correctly."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_boto.return_value = MagicMock()
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
assert client.bucket_name == "se3zdnb5o4"
|
|
assert client.endpoint_url == "https://s3api-eur-is-1.runpod.io"
|
|
mock_boto.assert_called_once()
|
|
|
|
def test_s3_client_initialization_failure(self, s3_config):
|
|
"""Test S3 client initialization failure."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_boto.side_effect = BotoCoreError()
|
|
|
|
with pytest.raises(S3Error, match="Failed to initialize"):
|
|
S3Client(**s3_config)
|
|
|
|
def test_upload_file_success(self, s3_config, temp_directory):
|
|
"""Test successful file upload."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
# Create test file
|
|
test_file = temp_directory / "test.bin"
|
|
test_file.write_bytes(b"test data")
|
|
|
|
result = client.upload_file(test_file, "binaries/test.bin")
|
|
|
|
assert result is True
|
|
mock_client.upload_file.assert_called_once()
|
|
|
|
def test_upload_file_not_found(self, s3_config, temp_directory):
|
|
"""Test upload with non-existent file."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_boto.return_value = MagicMock()
|
|
|
|
client = S3Client(**s3_config)
|
|
nonexistent = temp_directory / "nonexistent.bin"
|
|
|
|
with pytest.raises(FileNotFoundError):
|
|
client.upload_file(nonexistent, "binaries/test.bin")
|
|
|
|
def test_upload_file_s3_error(self, s3_config, temp_directory):
|
|
"""Test upload with S3 error."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.upload_file.side_effect = ClientError(
|
|
{"Error": {"Code": "NoSuchBucket", "Message": "Bucket not found"}},
|
|
"upload_file"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
test_file = temp_directory / "test.bin"
|
|
test_file.write_bytes(b"test data")
|
|
|
|
with pytest.raises(S3Error, match="Failed to upload"):
|
|
client.upload_file(test_file, "binaries/test.bin")
|
|
|
|
def test_upload_file_with_extra_args(self, s3_config, temp_directory):
|
|
"""Test upload with extra arguments."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
test_file = temp_directory / "test.bin"
|
|
test_file.write_bytes(b"test data")
|
|
|
|
extra_args = {"ContentType": "application/octet-stream"}
|
|
client.upload_file(test_file, "binaries/test.bin", extra_args)
|
|
|
|
call_args = mock_client.upload_file.call_args
|
|
assert call_args[1]["ExtraArgs"] == extra_args
|
|
|
|
def test_upload_binaries_success(self, s3_config, sample_binary_files):
|
|
"""Test uploading directory of binaries."""
|
|
binaries_dir, binary_files = sample_binary_files
|
|
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
uploaded = client.upload_binaries(binaries_dir, "binaries/")
|
|
|
|
# Should upload all binaries
|
|
assert len(uploaded) == 4
|
|
assert all(key.startswith("binaries/") for key in uploaded)
|
|
assert mock_client.upload_file.call_count == 4
|
|
|
|
def test_upload_binaries_not_directory(self, s3_config, temp_directory):
|
|
"""Test upload_binaries with non-directory."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_boto.return_value = MagicMock()
|
|
|
|
client = S3Client(**s3_config)
|
|
not_a_dir = temp_directory / "file.txt"
|
|
not_a_dir.write_text("test")
|
|
|
|
with pytest.raises(NotADirectoryError):
|
|
client.upload_binaries(not_a_dir)
|
|
|
|
def test_download_file_success(self, s3_config, temp_directory):
|
|
"""Test successful file download."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
local_path = temp_directory / "downloaded.bin"
|
|
|
|
result = client.download_file("models/test.safetensors", local_path)
|
|
|
|
assert result is True
|
|
mock_client.download_file.assert_called_once()
|
|
|
|
def test_download_file_creates_directories(self, s3_config, temp_directory):
|
|
"""Test download creates parent directories."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
local_path = temp_directory / "nested" / "dir" / "file.bin"
|
|
|
|
client.download_file("models/test.bin", local_path)
|
|
|
|
assert local_path.parent.exists()
|
|
|
|
def test_download_file_s3_error(self, s3_config, temp_directory):
|
|
"""Test download with S3 error."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.download_file.side_effect = ClientError(
|
|
{"Error": {"Code": "NoSuchKey", "Message": "Key not found"}},
|
|
"download_file"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
local_path = temp_directory / "test.bin"
|
|
|
|
with pytest.raises(S3Error, match="Failed to download"):
|
|
client.download_file("nonexistent.bin", local_path)
|
|
|
|
def test_stream_logs_from_start(self, s3_config, mock_log_content):
|
|
"""Test streaming logs from start."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_response = {
|
|
'Body': io.BytesIO(mock_log_content.encode('utf-8'))
|
|
}
|
|
mock_client.get_object.return_value = mock_response
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
content = client.stream_logs("logs/test.log", start_byte=0)
|
|
|
|
assert "Starting training" in content
|
|
mock_client.get_object.assert_called_once()
|
|
|
|
def test_stream_logs_with_offset(self, s3_config):
|
|
"""Test streaming logs from specific byte offset."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
partial_content = b"Partial log content"
|
|
mock_response = {
|
|
'Body': io.BytesIO(partial_content)
|
|
}
|
|
mock_client.get_object.return_value = mock_response
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
content = client.stream_logs("logs/test.log", start_byte=1000)
|
|
|
|
assert content == "Partial log content"
|
|
# Verify Range header was used
|
|
call_args = mock_client.get_object.call_args
|
|
assert call_args[1]["Range"] == "bytes=1000-"
|
|
|
|
def test_stream_logs_tail_bytes(self, s3_config):
|
|
"""Test streaming last N bytes of logs."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
|
|
# Mock head_object to return file size
|
|
mock_client.head_object.return_value = {'ContentLength': 5000}
|
|
|
|
# Mock get_object for range request
|
|
tail_content = b"Last 1KB of logs"
|
|
mock_response = {
|
|
'Body': io.BytesIO(tail_content)
|
|
}
|
|
mock_client.get_object.return_value = mock_response
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
content = client.stream_logs("logs/test.log", tail_bytes=1024)
|
|
|
|
assert content == "Last 1KB of logs"
|
|
# Should request from byte 3976 (5000 - 1024)
|
|
call_args = mock_client.get_object.call_args
|
|
assert call_args[1]["Range"] == "bytes=3976-"
|
|
|
|
def test_stream_logs_nonexistent_file(self, s3_config):
|
|
"""Test streaming logs from nonexistent file."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.get_object.side_effect = ClientError(
|
|
{"Error": {"Code": "NoSuchKey", "Message": "Key not found"}},
|
|
"get_object"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
content = client.stream_logs("logs/nonexistent.log")
|
|
|
|
# Should return empty string for nonexistent logs
|
|
assert content == ""
|
|
|
|
def test_stream_logs_s3_error(self, s3_config):
|
|
"""Test streaming logs with S3 error."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.get_object.side_effect = ClientError(
|
|
{"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
|
|
"get_object"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
with pytest.raises(S3Error, match="Failed to stream logs"):
|
|
client.stream_logs("logs/test.log")
|
|
|
|
def test_list_files_success(self, s3_config, mock_s3_file_list):
|
|
"""Test listing files in S3."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.return_value = mock_s3_file_list
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
files = client.list_files(prefix="models/")
|
|
|
|
assert len(files) == 3
|
|
assert files[0]["Key"] == "models/tft_best.safetensors"
|
|
assert files[0]["Size"] == 1024000
|
|
|
|
def test_list_files_empty(self, s3_config):
|
|
"""Test listing files with no results."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.return_value = {} # No Contents field
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
files = client.list_files(prefix="nonexistent/")
|
|
|
|
assert files == []
|
|
|
|
def test_list_files_with_max_keys(self, s3_config):
|
|
"""Test listing files with max_keys parameter."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.return_value = {"Contents": []}
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
client.list_files(prefix="models/", max_keys=500)
|
|
|
|
call_args = mock_client.list_objects_v2.call_args
|
|
assert call_args[1]["MaxKeys"] == 500
|
|
|
|
def test_list_files_s3_error(self, s3_config):
|
|
"""Test list_files with S3 error."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.side_effect = ClientError(
|
|
{"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
|
|
"list_objects_v2"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
with pytest.raises(S3Error, match="Failed to list files"):
|
|
client.list_files(prefix="models/")
|
|
|
|
def test_download_results_success(self, s3_config, temp_directory):
|
|
"""Test downloading training results."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
# Mock only returns model files (not logs)
|
|
mock_client.list_objects_v2.return_value = {
|
|
"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"
|
|
}
|
|
]
|
|
}
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
downloaded = client.download_results("models/", temp_directory)
|
|
|
|
# Should download 2 model files
|
|
assert len(downloaded) == 2
|
|
assert mock_client.download_file.call_count == 2
|
|
|
|
def test_download_results_with_patterns(self, s3_config, mock_s3_file_list, temp_directory):
|
|
"""Test downloading results with file patterns."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.return_value = mock_s3_file_list
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
downloaded = client.download_results(
|
|
"models/",
|
|
temp_directory,
|
|
file_patterns=["*.safetensors"]
|
|
)
|
|
|
|
# Should only download .safetensors files
|
|
assert len(downloaded) == 1
|
|
assert downloaded[0].name == "tft_best.safetensors"
|
|
|
|
def test_download_results_preserves_structure(self, s3_config, temp_directory):
|
|
"""Test download preserves directory structure."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.list_objects_v2.return_value = {
|
|
"Contents": [
|
|
{"Key": "models/subdir/file.bin", "Size": 1024, "LastModified": "2025-10-29"}
|
|
]
|
|
}
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
downloaded = client.download_results("models/", temp_directory)
|
|
|
|
# Should preserve subdirectory
|
|
assert len(downloaded) == 1
|
|
assert "subdir" in str(downloaded[0])
|
|
|
|
def test_delete_file_success(self, s3_config):
|
|
"""Test deleting a file."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
result = client.delete_file("models/old_model.bin")
|
|
|
|
assert result is True
|
|
mock_client.delete_object.assert_called_once()
|
|
|
|
def test_delete_file_error(self, s3_config):
|
|
"""Test delete with S3 error."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
mock_client.delete_object.side_effect = ClientError(
|
|
{"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
|
|
"delete_object"
|
|
)
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
|
|
with pytest.raises(S3Error, match="Failed to delete"):
|
|
client.delete_file("models/test.bin")
|
|
|
|
@pytest.mark.parametrize("prefix,expected_count", [
|
|
("models/", 2), # 2 model files
|
|
("logs/", 1), # 1 log file
|
|
("", 3), # All files
|
|
])
|
|
def test_list_files_with_different_prefixes(self, s3_config, mock_s3_file_list, prefix, expected_count):
|
|
"""Test listing files with different prefixes."""
|
|
with patch('boto3.client') as mock_boto:
|
|
mock_client = MagicMock()
|
|
|
|
# Filter mock data based on prefix
|
|
filtered_contents = [
|
|
item for item in mock_s3_file_list["Contents"]
|
|
if item["Key"].startswith(prefix)
|
|
]
|
|
mock_client.list_objects_v2.return_value = {"Contents": filtered_contents}
|
|
mock_boto.return_value = mock_client
|
|
|
|
client = S3Client(**s3_config)
|
|
files = client.list_files(prefix=prefix)
|
|
|
|
assert len(files) == expected_count
|