File size: 4,045 Bytes
ce4bc73 |
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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
"""Tests for the Gemini AI client."""
import os
import unittest
from unittest.mock import MagicMock, patch
import pytest
from src.folio.gemini_client import GeminiClient
class TestGeminiClient(unittest.TestCase):
"""Tests for the Gemini AI client."""
def setUp(self):
"""Set up test environment."""
# Save original environment
self.original_env = os.environ.copy()
# Set test API key
os.environ["GEMINI_API_KEY"] = "test_api_key"
def tearDown(self):
"""Clean up test environment."""
# Restore original environment
os.environ.clear()
os.environ.update(self.original_env)
@patch("google.generativeai.configure")
@patch("google.generativeai.GenerativeModel")
def test_init_with_api_key(self, mock_model, mock_configure):
"""Test that the client initializes correctly with an API key."""
# Set up mock
mock_model.return_value = MagicMock()
# Initialize client
client = GeminiClient()
# Check that configure was called with the API key
mock_configure.assert_called_once_with(api_key="test_api_key")
# Check that the model was initialized
mock_model.assert_called_once()
# Check that the model is available
assert hasattr(client, "model")
@patch("google.generativeai.configure")
@patch("google.generativeai.GenerativeModel")
def test_init_without_api_key(self, mock_model, mock_configure):
"""Test that the client raises an error when no API key is provided."""
# Remove API key from environment
if "GEMINI_API_KEY" in os.environ:
del os.environ["GEMINI_API_KEY"]
# Check that initializing the client raises an error
with pytest.raises(
ValueError, match="GEMINI_API_KEY environment variable not set"
):
GeminiClient()
# Check that configure was not called
mock_configure.assert_not_called()
# Check that the model was not initialized
mock_model.assert_not_called()
@patch("google.generativeai.configure")
@patch("google.generativeai.GenerativeModel")
def test_chat_sync_basic(self, mock_model, _mock_configure):
"""Test that the chat_sync method works correctly."""
# Set up mocks
mock_model_instance = MagicMock()
mock_model.return_value = mock_model_instance
# Set up chat session mock
mock_chat = MagicMock()
mock_model_instance.start_chat.return_value = mock_chat
# Set up response mock
mock_response = MagicMock()
mock_response.text = "Test response"
mock_chat.send_message.return_value = mock_response
# Initialize client
client = GeminiClient()
# Call chat_sync
result = client.chat_sync("Test message", [])
# Check that the chat session was created
mock_model_instance.start_chat.assert_called_once()
# Check that send_message was called with the message
mock_chat.send_message.assert_called_once()
# Check the result
assert result["response"] == "Test response"
assert result["complete"] is True
@patch("google.generativeai.configure")
@patch("google.generativeai.GenerativeModel")
def test_chat_sync_with_error(self, mock_model, _mock_configure):
"""Test that the chat_sync method handles errors correctly."""
# Set up mocks
mock_model_instance = MagicMock()
mock_model.return_value = mock_model_instance
# Set up chat session mock to raise an exception
mock_model_instance.start_chat.side_effect = Exception("Test error")
# Initialize client
client = GeminiClient()
# Call chat_sync
result = client.chat_sync("Test message", [])
# Check the result
assert "error" in result
assert result["error"] is True
assert "Test error" in result["response"]
assert result["complete"] is False
|