File size: 3,069 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 |
"""
Test filename uniqueness for generate_image
"""
import os
import sys
from unittest.mock import Mock, patch
from dotenv import load_dotenv
import time
# Load environment variables
load_dotenv()
# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
from agent import VisualizationAgent
from PIL import Image
def test_filename_uniqueness():
"""Test that multiple image generations create unique filenames"""
print("\n=== Testing Filename Uniqueness ===")
try:
# Create a simple test image
test_image = Image.new('RGB', (100, 100), color='green')
# Create mock part with inline_data
mock_part = Mock()
mock_part.inline_data = Mock()
mock_part.as_image = Mock(return_value=test_image)
# Create mock response
mock_response = Mock()
mock_response.parts = [mock_part]
with patch('agent.genai.Client') as MockClient:
mock_client_instance = Mock()
mock_client_instance.models.generate_content.return_value = mock_response
MockClient.return_value = mock_client_instance
agent = VisualizationAgent(api_key="test_key")
# Generate multiple images
filenames = []
for i in range(5):
result = agent.generate_image(f"Test prompt {i}")
assert result["success"] is True, f"Generation {i} failed"
filenames.append(result["image_path"])
# Small delay to ensure timestamp changes
time.sleep(0.01)
# Verify all filenames are unique
unique_filenames = set(filenames)
assert len(unique_filenames) == len(filenames), \
f"Duplicate filenames found! {len(unique_filenames)} unique out of {len(filenames)}"
print("✓ All filenames are unique")
print(f" Generated {len(filenames)} images with unique names")
# Verify all files exist
for filepath in filenames:
assert os.path.exists(filepath), f"File not found: {filepath}"
print("✓ All image files created successfully")
# Clean up
for filepath in filenames:
if os.path.exists(filepath):
os.remove(filepath)
print(" Cleaned up test images")
return True
except Exception as e:
print(f"✗ Test failed: {e}")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
print("=" * 70)
print("Testing Filename Uniqueness")
print("=" * 70)
if test_filename_uniqueness():
print("\n✓ All tests passed")
sys.exit(0)
else:
print("\n✗ Tests failed")
sys.exit(1)
|