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