File size: 3,792 Bytes
9b24c4d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
"""
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")
|