visualization-agent/test_file_management.py

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