Spaces:
Sleeping
Sleeping
File size: 3,574 Bytes
235461a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from unittest.mock import Mock, patch
import pytest
# Mock dependencies before importing
with patch("langgraph.checkpoint.postgres.PostgresSaver"):
with patch("llm.connection_manager.get_checkpointer"):
with patch("llm.connection_manager._test_connection", return_value=True):
from llm.utils import upload_generated_image_to_s3
class TestS3Utils:
"""Test cases for S3 utility functions."""
@patch("llm.utils.boto3.client")
@patch.dict("os.environ", {"AWS_S3_BUCKET_NAME": "test-bucket"})
def test_upload_generated_image_success(self, mock_boto3_client):
"""Test successful image upload to S3."""
mock_s3_client = Mock()
mock_boto3_client.return_value = mock_s3_client
mock_s3_client.generate_presigned_url.return_value = "https://test-bucket.s3.amazonaws.com/test-url"
mock_s3_client.head_object.return_value = {
"Metadata": {
"title": "Test Image",
"imageId": "test-uuid-123",
"userId": "test_user",
"uploadedAt": "2024-01-01T00:00:00Z",
"type": "generated",
"generationPrompt": "A beautiful sunset",
}
}
result = upload_generated_image_to_s3(b"fake_image_data", "test-uuid-123", "test_user", "A beautiful sunset", "Test Image")
assert result["success"] is True
assert result["url"] == "https://test-bucket.s3.amazonaws.com/test-url"
assert result["image_id"] == "test-uuid-123"
# Verify S3 calls
mock_s3_client.put_object.assert_called_once()
put_call = mock_s3_client.put_object.call_args
assert put_call[1]["Bucket"] == "test-bucket"
assert put_call[1]["Key"] == "users/test_user/images/test-uuid-123"
assert put_call[1]["Body"] == b"fake_image_data"
assert put_call[1]["ContentType"] == "image/png"
@patch("llm.utils.boto3.client")
def test_upload_generated_image_missing_bucket(self, mock_boto3_client):
"""Test upload when S3 bucket name is not set."""
with patch.dict("os.environ", {}, clear=True):
result = upload_generated_image_to_s3(b"data", "id", "user", "prompt", "title")
assert result["success"] is False
assert "AWS_S3_BUCKET_NAME environment variable is not set" in result["error"]
@patch("llm.utils.boto3.client")
def test_upload_generated_image_s3_error(self, mock_boto3_client):
"""Test upload when S3 operations fail."""
mock_s3_client = Mock()
mock_s3_client.put_object.side_effect = Exception("S3 upload failed")
mock_boto3_client.return_value = mock_s3_client
result = upload_generated_image_to_s3(b"data", "id", "user", "prompt", "title")
assert result["success"] is False
assert "S3 upload failed" in result["error"]
@patch("llm.utils.boto3.client")
def test_upload_generated_image_default_title(self, mock_boto3_client):
"""Test upload with default title when none provided."""
mock_s3_client = Mock()
mock_boto3_client.return_value = mock_s3_client
mock_s3_client.generate_presigned_url.return_value = "https://test-url"
mock_s3_client.head_object.return_value = {"Metadata": {}}
upload_generated_image_to_s3(b"data", "id", "user", "prompt")
put_call = mock_s3_client.put_object.call_args
metadata = put_call[1]["Metadata"]
assert metadata["title"] == "Generated Image"
if __name__ == "__main__":
pytest.main([__file__])
|