visualization-agent/test_env_config.py

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