""" 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)