File size: 7,764 Bytes
9b24c4d |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
"""
Tests for file management functionality (Task 9)
Tests Requirements 3.2, 3.3, 3.4
"""
import os
import base64
import tempfile
import shutil
from unittest.mock import Mock, patch, MagicMock
from PIL import Image
from io import BytesIO
from agent import VisualizationAgent
def test_save_image_creates_png_file():
"""Test that _save_image creates a valid PNG file"""
# Create a test image
test_image = Image.new('RGB', (100, 100), color='red')
# Create temporary directory
with tempfile.TemporaryDirectory() as temp_dir:
agent = VisualizationAgent(api_key="test_key", output_dir=temp_dir)
# Save image
filename = "test_image.png"
filepath = agent._save_image(test_image, filename)
# Verify file exists
assert os.path.exists(filepath)
# Verify it's a PNG file
assert filepath.endswith('.png')
# Verify we can open it as an image
loaded_image = Image.open(filepath)
assert loaded_image.size == (100, 100)
def test_save_image_validates_directory_writable():
"""Test that _save_image checks if directory is writable"""
test_image = Image.new('RGB', (100, 100), color='blue')
with tempfile.TemporaryDirectory() as temp_dir:
# Create a read-only directory
readonly_dir = os.path.join(temp_dir, "readonly")
os.makedirs(readonly_dir)
os.chmod(readonly_dir, 0o444) # Read-only
try:
agent = VisualizationAgent(api_key="test_key", output_dir=readonly_dir)
# Attempt to save should raise an exception
try:
agent._save_image(test_image, "test.png")
assert False, "Should have raised an exception for read-only directory"
except Exception as e:
assert "Permission denied" in str(e) or "not writable" in str(e)
finally:
# Restore permissions for cleanup
os.chmod(readonly_dir, 0o755)
def test_unique_filename_generation():
"""Test that filenames are unique using timestamp and UUID"""
with tempfile.TemporaryDirectory() as temp_dir:
agent = VisualizationAgent(api_key="test_key", output_dir=temp_dir)
# Mock the API response
mock_response = Mock()
mock_part = Mock()
mock_image = Image.new('RGB', (100, 100), color='green')
mock_part.inline_data = True
mock_part.as_image.return_value = mock_image
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content', return_value=mock_response):
# Generate multiple images
result1 = agent.generate_image("test prompt 1")
result2 = agent.generate_image("test prompt 2")
# Verify both succeeded
assert result1['success']
assert result2['success']
# Verify filenames are different
path1 = result1.get('image_path')
path2 = result2.get('image_path')
if path1 and path2: # Only check if both paths exist (may be None in stateless env)
assert path1 != path2
assert os.path.basename(path1) != os.path.basename(path2)
def test_image_to_base64_encoding():
"""Test that _image_to_base64 correctly encodes images"""
test_image = Image.new('RGB', (50, 50), color='yellow')
agent = VisualizationAgent(api_key="test_key")
# Encode to base64
base64_str = agent._image_to_base64(test_image)
# Verify it's a valid base64 string
assert isinstance(base64_str, str)
assert len(base64_str) > 0
# Verify we can decode it back to an image
image_bytes = base64.b64decode(base64_str)
decoded_image = Image.open(BytesIO(image_bytes))
assert decoded_image.size == (50, 50)
assert decoded_image.mode == 'RGB'
def test_generate_image_returns_base64():
"""Test that generate_image returns base64 encoded image for stateless deployment"""
agent = VisualizationAgent(api_key="test_key")
# Mock the API response
mock_response = Mock()
mock_part = Mock()
mock_image = Image.new('RGB', (100, 100), color='purple')
mock_part.inline_data = True
mock_part.as_image.return_value = mock_image
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content', return_value=mock_response):
result = agent.generate_image("test prompt")
# Verify success
assert result['success']
# Verify base64 is present
assert 'image_base64' in result
assert result['image_base64'] is not None
assert len(result['image_base64']) > 0
# Verify we can decode the base64
image_bytes = base64.b64decode(result['image_base64'])
decoded_image = Image.open(BytesIO(image_bytes))
assert decoded_image.size == (100, 100)
def test_generate_image_handles_save_failure_gracefully():
"""Test that generate_image continues even if file save fails (stateless mode)"""
# Create a temporary directory that we'll make read-only
with tempfile.TemporaryDirectory() as temp_dir:
readonly_dir = os.path.join(temp_dir, "readonly")
os.makedirs(readonly_dir)
agent = VisualizationAgent(api_key="test_key", output_dir=readonly_dir)
# Make directory read-only after agent initialization
os.chmod(readonly_dir, 0o444)
try:
# Mock the API response
mock_response = Mock()
mock_part = Mock()
mock_image = Image.new('RGB', (100, 100), color='orange')
mock_part.inline_data = True
mock_part.as_image.return_value = mock_image
mock_response.parts = [mock_part]
with patch.object(agent.client.models, 'generate_content', return_value=mock_response):
result = agent.generate_image("test prompt")
# Should still succeed because base64 encoding works
assert result['success']
# Base64 should be present
assert result['image_base64'] is not None
# image_path may be None (save failed) but that's okay
# The important thing is the request didn't fail
finally:
# Restore permissions for cleanup
os.chmod(readonly_dir, 0o755)
def test_output_directory_creation():
"""Test that output directory is created if it doesn't exist"""
with tempfile.TemporaryDirectory() as temp_dir:
# Create a path that doesn't exist yet
new_dir = os.path.join(temp_dir, "new_output_dir")
assert not os.path.exists(new_dir)
# Initialize agent with non-existent directory
agent = VisualizationAgent(api_key="test_key", output_dir=new_dir)
# Verify directory was created
assert os.path.exists(new_dir)
assert os.path.isdir(new_dir)
def test_png_format_compliance():
"""Test that all saved images are in PNG format"""
test_image = Image.new('RGB', (100, 100), color='cyan')
with tempfile.TemporaryDirectory() as temp_dir:
agent = VisualizationAgent(api_key="test_key", output_dir=temp_dir)
# Save image
filepath = agent._save_image(test_image, "test.png")
# Open and verify format
loaded_image = Image.open(filepath)
assert loaded_image.format == 'PNG'
if __name__ == "__main__":
import pytest
pytest.main([__file__, "-v"])
|