""" Test environment variable configuration for VisualizationAgent """ import os import sys from unittest.mock import 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__))) from agent import VisualizationAgent def test_visualization_model_env_var(): """Test that VISUALIZATION_MODEL environment variable is used""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key', 'VISUALIZATION_MODEL': 'gemini-3-pro-image-preview' }): agent = VisualizationAgent() assert agent.model == 'gemini-3-pro-image-preview' print("✓ VISUALIZATION_MODEL environment variable works correctly") def test_visualization_output_dir_env_var(): """Test that VISUALIZATION_OUTPUT_DIR environment variable is used""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key', 'VISUALIZATION_OUTPUT_DIR': './custom_output' }): agent = VisualizationAgent() assert agent.output_dir == './custom_output' print("✓ VISUALIZATION_OUTPUT_DIR environment variable works correctly") def test_env_var_priority_model(): """Test that constructor parameter takes priority over environment variable for model""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key', 'VISUALIZATION_MODEL': 'gemini-3-pro-image-preview' }): agent = VisualizationAgent(model='gemini-2.5-flash-image') assert agent.model == 'gemini-2.5-flash-image' print("✓ Constructor parameter takes priority over VISUALIZATION_MODEL env var") def test_env_var_priority_output_dir(): """Test that constructor parameter takes priority over environment variable for output_dir""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key', 'VISUALIZATION_OUTPUT_DIR': './env_output' }): agent = VisualizationAgent(output_dir='./param_output') assert agent.output_dir == './param_output' print("✓ Constructor parameter takes priority over VISUALIZATION_OUTPUT_DIR env var") def test_default_values_when_no_env_vars(): """Test that default values are used when no environment variables are set""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key' }, clear=True): # Clear VISUALIZATION_MODEL and VISUALIZATION_OUTPUT_DIR os.environ.pop('VISUALIZATION_MODEL', None) os.environ.pop('VISUALIZATION_OUTPUT_DIR', None) agent = VisualizationAgent() assert agent.model == 'gemini-2.5-flash-image' assert agent.output_dir == './generated_images' print("✓ Default values used when environment variables not set") def test_all_env_vars_together(): """Test that all environment variables work together""" with patch.dict(os.environ, { 'GEMINI_API_KEY': 'test_key_123', 'VISUALIZATION_MODEL': 'gemini-3-pro-image-preview', 'VISUALIZATION_OUTPUT_DIR': './test_images' }): agent = VisualizationAgent() assert agent.api_key == 'test_key_123' assert agent.model == 'gemini-3-pro-image-preview' assert agent.output_dir == './test_images' print("✓ All environment variables work together correctly") if __name__ == '__main__': print("\nTesting Environment Variable Configuration\n") print("=" * 60) test_visualization_model_env_var() test_visualization_output_dir_env_var() test_env_var_priority_model() test_env_var_priority_output_dir() test_default_values_when_no_env_vars() test_all_env_vars_together() print("=" * 60) print("\n✅ All environment variable configuration tests passed!\n")