Ed_
8aa70e287f
This change corrects the implementation of get_gemini_cache_stats to use the Gemini client instance and updates the corresponding test to use proper mocking.
46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Import the necessary functions from ai_client, including the reset helper
|
|
from ai_client import get_gemini_cache_stats, reset_session
|
|
|
|
def test_get_gemini_cache_stats_with_mock_client():
|
|
"""
|
|
Test that get_gemini_cache_stats correctly processes cache lists
|
|
from a mocked client instance.
|
|
"""
|
|
# Ensure a clean state before the test by resetting the session
|
|
reset_session()
|
|
|
|
# 1. Create a mock for the cache object that the client will return
|
|
mock_cache = MagicMock()
|
|
mock_cache.name = "cachedContents/test-cache"
|
|
mock_cache.display_name = "Test Cache"
|
|
mock_cache.model = "models/gemini-1.5-pro-001"
|
|
mock_cache.size_bytes = 1024
|
|
|
|
# 2. Create a mock for the client instance
|
|
mock_client_instance = MagicMock()
|
|
# Configure its `caches.list` method to return our mock cache
|
|
mock_client_instance.caches.list.return_value = [mock_cache]
|
|
|
|
# 3. Patch the Client constructor to return our mock instance
|
|
# This intercepts the `_ensure_gemini_client` call inside the function
|
|
with patch('google.genai.Client', return_value=mock_client_instance) as mock_client_constructor:
|
|
|
|
# 4. Call the function under test
|
|
stats = get_gemini_cache_stats()
|
|
|
|
# 5. Assert that the function behaved as expected
|
|
|
|
# It should have constructed the client
|
|
mock_client_constructor.assert_called_once()
|
|
# It should have called the `list` method on the `caches` attribute
|
|
mock_client_instance.caches.list.assert_called_once()
|
|
|
|
# The returned stats dictionary should be correct
|
|
assert "cache_count" in stats
|
|
assert "total_size_bytes" in stats
|
|
assert stats["cache_count"] == 1
|
|
assert stats["total_size_bytes"] == 1024
|