visualization-agent/test_filename_uniqueness.py

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