""" Simple test for generate_image method """ import os import sys from unittest.mock import Mock, patch from dotenv import load_dotenv # 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 import io def test_generate_image_with_mock(): """Test generate_image method with mocked API response""" print("\n=== Testing generate_image Method ===") try: # Create a simple test image test_image = Image.new('RGB', (100, 100), color='red') # 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] # Mock the genai.Client 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 # Create agent agent = VisualizationAgent(api_key="test_key") # Test generate_image result = agent.generate_image("Test prompt for disaster-resistant building") # Verify result assert result["success"] is True, f"Expected success=True, got {result}" assert "image_path" in result, "Missing image_path in result" assert result["image_path"] is not None, "image_path is None" assert os.path.exists(result["image_path"]), f"Image file not created: {result['image_path']}" print("✓ generate_image method works correctly") print(f" Image saved to: {result['image_path']}") # Clean up if os.path.exists(result["image_path"]): os.remove(result["image_path"]) print(" Cleaned up test image") return True except Exception as e: print(f"✗ Test failed: {e}") import traceback traceback.print_exc() return False def test_generate_image_empty_prompt(): """Test generate_image with empty prompt""" print("\n=== Testing generate_image with Empty Prompt ===") try: with patch('agent.genai.Client') as MockClient: mock_client_instance = Mock() MockClient.return_value = mock_client_instance agent = VisualizationAgent(api_key="test_key") # Test with empty prompt result = agent.generate_image("") assert result["success"] is False, "Should fail with empty prompt" assert "error" in result, "Missing error field" assert result["error"]["code"] == "INVALID_INPUT", f"Wrong error code: {result['error']['code']}" print("✓ Empty prompt validation works") print(f" Error: {result['error']['message']}") return True except Exception as e: print(f"✗ Test failed: {e}") import traceback traceback.print_exc() return False def test_generate_image_no_image_data(): """Test generate_image when API returns no image data""" print("\n=== Testing generate_image with No Image Data ===") try: # Create mock response with no inline_data mock_part = Mock() mock_part.inline_data = None 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") result = agent.generate_image("Test prompt") assert result["success"] is False, "Should fail when no image data" assert "error" in result, "Missing error field" assert result["error"]["code"] == "NO_IMAGE_DATA", f"Wrong error code: {result['error']['code']}" print("✓ No image data handling works") print(f" Error: {result['error']['message']}") return True except Exception as e: print(f"✗ Test failed: {e}") import traceback traceback.print_exc() return False def test_generate_image_with_config(): """Test generate_image with custom config""" print("\n=== Testing generate_image with Custom Config ===") try: # Create a simple test image test_image = Image.new('RGB', (100, 100), color='blue') # 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") # Test with custom config config = { "model": "gemini-3-pro-image-preview", "aspect_ratio": "16:9" } result = agent.generate_image("Test prompt", config=config) assert result["success"] is True, f"Expected success=True, got {result}" # Verify the model was passed correctly call_args = mock_client_instance.models.generate_content.call_args assert call_args[1]["model"] == "gemini-3-pro-image-preview", "Model not passed correctly" print("✓ Custom config handling works") print(f" Model used: {config['model']}") # Clean up if result.get("image_path") and os.path.exists(result["image_path"]): os.remove(result["image_path"]) 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 generate_image Implementation") print("=" * 70) tests = [ test_generate_image_with_mock, test_generate_image_empty_prompt, test_generate_image_no_image_data, test_generate_image_with_config, ] passed = 0 failed = 0 for test in tests: try: if test(): passed += 1 else: failed += 1 except Exception as e: print(f"✗ Test crashed: {e}") failed += 1 print("\n" + "=" * 70) print(f"Results: {passed} passed, {failed} failed") print("=" * 70) sys.exit(0 if failed == 0 else 1)