83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
# be/tests/core/test_gemini.py
|
|
import pytest
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Store original key if exists, then clear it for testing missing key scenario
|
|
original_api_key = os.environ.get("GEMINI_API_KEY")
|
|
if "GEMINI_API_KEY" in os.environ:
|
|
del os.environ["GEMINI_API_KEY"]
|
|
|
|
# --- Test Module Import ---
|
|
# This forces the module-level initialization code in gemini.py to run
|
|
# We need to reload modules because settings might have been cached
|
|
from importlib import reload
|
|
from app.config import settings as app_settings
|
|
from app.core import gemini as gemini_core
|
|
|
|
# Reload settings first to ensure GEMINI_API_KEY is None initially
|
|
reload(app_settings)
|
|
# Reload gemini core to trigger initialization logic with potentially missing key
|
|
reload(gemini_core)
|
|
|
|
|
|
def test_gemini_initialization_without_key():
|
|
"""Verify behavior when GEMINI_API_KEY is not set."""
|
|
# Reload modules again to ensure clean state for this specific test
|
|
if "GEMINI_API_KEY" in os.environ:
|
|
del os.environ["GEMINI_API_KEY"]
|
|
reload(app_settings)
|
|
reload(gemini_core)
|
|
|
|
assert gemini_core.gemini_flash_client is None
|
|
assert gemini_core.gemini_initialization_error is not None
|
|
assert "GEMINI_API_KEY not configured" in gemini_core.gemini_initialization_error
|
|
|
|
with pytest.raises(RuntimeError, match="GEMINI_API_KEY not configured"):
|
|
gemini_core.get_gemini_client()
|
|
|
|
@patch('google.generativeai.configure')
|
|
@patch('google.generativeai.GenerativeModel')
|
|
def test_gemini_initialization_with_key(mock_generative_model: MagicMock, mock_configure: MagicMock):
|
|
"""Verify initialization logic is called when key is present (using mocks)."""
|
|
# Set a dummy key in the environment for this test
|
|
test_key = "TEST_API_KEY_123"
|
|
os.environ["GEMINI_API_KEY"] = test_key
|
|
|
|
# Reload settings and gemini module to pick up the new key
|
|
reload(app_settings)
|
|
reload(gemini_core)
|
|
|
|
# Assertions
|
|
mock_configure.assert_called_once_with(api_key=test_key)
|
|
mock_generative_model.assert_called_once_with(
|
|
model_name="gemini-1.5-flash-latest",
|
|
safety_settings=pytest.ANY, # Check safety settings were passed (ANY allows flexibility)
|
|
# generation_config=pytest.ANY # Check if you added default generation config
|
|
)
|
|
assert gemini_core.gemini_flash_client is not None
|
|
assert gemini_core.gemini_initialization_error is None
|
|
|
|
# Test get_gemini_client() success path
|
|
client = gemini_core.get_gemini_client()
|
|
assert client is not None # Should return the mocked client instance
|
|
|
|
# Clean up environment variable after test
|
|
if original_api_key:
|
|
os.environ["GEMINI_API_KEY"] = original_api_key
|
|
else:
|
|
if "GEMINI_API_KEY" in os.environ:
|
|
del os.environ["GEMINI_API_KEY"]
|
|
# Reload modules one last time to restore state for other tests
|
|
reload(app_settings)
|
|
reload(gemini_core)
|
|
|
|
# Restore original key after all tests in the module run (if needed)
|
|
def teardown_module(module):
|
|
if original_api_key:
|
|
os.environ["GEMINI_API_KEY"] = original_api_key
|
|
else:
|
|
if "GEMINI_API_KEY" in os.environ:
|
|
del os.environ["GEMINI_API_KEY"]
|
|
reload(app_settings)
|
|
reload(gemini_core) |