mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 23:07:47 +00:00
* add vertex tests * add uperbound * add pagination tests
This commit is contained in:
Aarish Alam
committed by
Sameer Kankute
Sameer Kankute
parent
b3f1696946
commit
3e04e20205
@@ -27,6 +27,8 @@ local_cache_obj = Cache(
|
||||
type=LiteLLMCacheType.LOCAL
|
||||
) # only used for calling 'get_cache_key' function
|
||||
|
||||
MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination
|
||||
|
||||
|
||||
class ContextCachingEndpoints(VertexBase):
|
||||
"""
|
||||
@@ -115,7 +117,7 @@ class ContextCachingEndpoints(VertexBase):
|
||||
- None
|
||||
"""
|
||||
|
||||
_, url = self._get_token_and_url_context_caching(
|
||||
_, base_url = self._get_token_and_url_context_caching(
|
||||
gemini_api_key=api_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
@@ -123,43 +125,63 @@ class ContextCachingEndpoints(VertexBase):
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=vertex_auth_header
|
||||
)
|
||||
try:
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {},
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
resp = client.get(url=url, headers=headers)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
page_token: Optional[str] = None
|
||||
|
||||
# Iterate through all pages
|
||||
for _ in range(MAX_PAGINATION_PAGES):
|
||||
# Build URL with pagination token if present
|
||||
if page_token:
|
||||
separator = "&" if "?" in base_url else "?"
|
||||
url = f"{base_url}{separator}pageToken={page_token}"
|
||||
else:
|
||||
url = base_url
|
||||
|
||||
try:
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {},
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
resp = client.get(url=url, headers=headers)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
raw_response = resp.json()
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
raw_response = resp.json()
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
# Check current page for matching cache_key
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
# Check if there are more pages
|
||||
page_token = all_cached_items.get("nextPageToken")
|
||||
if not page_token:
|
||||
# No more pages, cache not found
|
||||
break
|
||||
|
||||
return None
|
||||
|
||||
@@ -187,7 +209,7 @@ class ContextCachingEndpoints(VertexBase):
|
||||
- None
|
||||
"""
|
||||
|
||||
_, url = self._get_token_and_url_context_caching(
|
||||
_, base_url = self._get_token_and_url_context_caching(
|
||||
gemini_api_key=api_key,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
@@ -195,43 +217,63 @@ class ContextCachingEndpoints(VertexBase):
|
||||
vertex_location=vertex_location,
|
||||
vertex_auth_header=vertex_auth_header
|
||||
)
|
||||
try:
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {},
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
resp = await client.get(url=url, headers=headers)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
page_token: Optional[str] = None
|
||||
|
||||
# Iterate through all pages
|
||||
for _ in range(MAX_PAGINATION_PAGES):
|
||||
# Build URL with pagination token if present
|
||||
if page_token:
|
||||
separator = "&" if "?" in base_url else "?"
|
||||
url = f"{base_url}{separator}pageToken={page_token}"
|
||||
else:
|
||||
url = base_url
|
||||
|
||||
try:
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": {},
|
||||
"api_base": url,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
resp = await client.get(url=url, headers=headers)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 403:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
raw_response = resp.json()
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code, message=e.response.text
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
raw_response = resp.json()
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
# Check current page for matching cache_key
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
# Check if there are more pages
|
||||
page_token = all_cached_items.get("nextPageToken")
|
||||
if not page_token:
|
||||
# No more pages, cache not found
|
||||
break
|
||||
|
||||
return None
|
||||
|
||||
@@ -501,4 +543,4 @@ class ContextCachingEndpoints(VertexBase):
|
||||
pass
|
||||
|
||||
async def async_get_cache(self):
|
||||
pass
|
||||
pass
|
||||
@@ -14,6 +14,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
|
||||
MAX_PAGINATION_PAGES,
|
||||
ContextCachingEndpoints,
|
||||
)
|
||||
|
||||
@@ -787,6 +788,358 @@ class TestContextCachingEndpoints:
|
||||
assert original_tools == self.sample_tools
|
||||
|
||||
|
||||
class TestCheckCachePagination:
|
||||
"""Test pagination logic in check_cache and async_check_cache methods."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Setup for each test method"""
|
||||
self.context_caching = ContextCachingEndpoints()
|
||||
self.mock_logging = MagicMock(spec=Logging)
|
||||
self.mock_client = MagicMock(spec=HTTPHandler)
|
||||
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_finds_cache_on_second_page(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that check_cache correctly handles pagination and finds cache on second page"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "target_cache_key"
|
||||
|
||||
# Mock first page response (no match, has nextPageToken)
|
||||
first_page_response = MagicMock()
|
||||
first_page_response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_1", "displayName": "cache_key_1"},
|
||||
{"name": "cache_2", "displayName": "cache_key_2"},
|
||||
],
|
||||
"nextPageToken": "token_page_2",
|
||||
}
|
||||
|
||||
# Mock second page response (has match, no nextPageToken)
|
||||
second_page_response = MagicMock()
|
||||
second_page_response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_3", "displayName": cache_key_to_find},
|
||||
{"name": "cache_4", "displayName": "cache_key_4"},
|
||||
]
|
||||
}
|
||||
|
||||
# Setup mock client to return different responses
|
||||
self.mock_client.get.side_effect = [first_page_response, second_page_response]
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result == "cache_3"
|
||||
assert self.mock_client.get.call_count == 2
|
||||
# Check that second call includes pageToken
|
||||
second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"]
|
||||
assert "pageToken=token_page_2" in second_call_url
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_stops_when_no_next_token(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that check_cache stops pagination when no nextPageToken is present"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Mock response without nextPageToken
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_1", "displayName": "cache_key_1"},
|
||||
{"name": "cache_2", "displayName": "cache_key_2"},
|
||||
]
|
||||
}
|
||||
|
||||
self.mock_client.get.return_value = response
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
assert self.mock_client.get.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_multiple_pages(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that check_cache correctly iterates through multiple pages"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "target_cache_key"
|
||||
|
||||
# Mock three pages
|
||||
page1 = MagicMock()
|
||||
page1.json.return_value = {
|
||||
"cachedContents": [{"name": "cache_1", "displayName": "cache_key_1"}],
|
||||
"nextPageToken": "token_page_2",
|
||||
}
|
||||
|
||||
page2 = MagicMock()
|
||||
page2.json.return_value = {
|
||||
"cachedContents": [{"name": "cache_2", "displayName": "cache_key_2"}],
|
||||
"nextPageToken": "token_page_3",
|
||||
}
|
||||
|
||||
page3 = MagicMock()
|
||||
page3.json.return_value = {
|
||||
"cachedContents": [{"name": "cache_3", "displayName": cache_key_to_find}],
|
||||
}
|
||||
|
||||
self.mock_client.get.side_effect = [page1, page2, page3]
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result == "cache_3"
|
||||
assert self.mock_client.get.call_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_finds_cache_on_second_page(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that async_check_cache correctly handles pagination and finds cache on second page"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "target_cache_key"
|
||||
|
||||
# Mock first page response (no match, has nextPageToken)
|
||||
first_page_response = MagicMock()
|
||||
first_page_response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_1", "displayName": "cache_key_1"},
|
||||
{"name": "cache_2", "displayName": "cache_key_2"},
|
||||
],
|
||||
"nextPageToken": "token_page_2",
|
||||
}
|
||||
|
||||
# Mock second page response (has match, no nextPageToken)
|
||||
second_page_response = MagicMock()
|
||||
second_page_response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_3", "displayName": cache_key_to_find},
|
||||
{"name": "cache_4", "displayName": "cache_key_4"},
|
||||
]
|
||||
}
|
||||
|
||||
# Setup mock async client to return different responses
|
||||
self.mock_async_client.get = AsyncMock(
|
||||
side_effect=[first_page_response, second_page_response]
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_async_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result == "cache_3"
|
||||
assert self.mock_async_client.get.call_count == 2
|
||||
# Check that second call includes pageToken
|
||||
second_call_url = self.mock_async_client.get.call_args_list[1].kwargs["url"]
|
||||
assert "pageToken=token_page_2" in second_call_url
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_stops_when_no_next_token(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that async_check_cache stops pagination when no nextPageToken is present"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Mock response without nextPageToken
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": "cache_1", "displayName": "cache_key_1"},
|
||||
{"name": "cache_2", "displayName": "cache_key_2"},
|
||||
]
|
||||
}
|
||||
|
||||
self.mock_async_client.get = AsyncMock(return_value=response)
|
||||
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_async_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is None
|
||||
assert self.mock_async_client.get.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
def test_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Create mock response that always has nextPageToken (infinite pagination scenario)
|
||||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
||||
# Create MAX_PAGINATION_PAGES responses, each with a nextPageToken
|
||||
self.mock_client.get.side_effect = [
|
||||
create_page_response(i) for i in range(MAX_PAGINATION_PAGES)
|
||||
]
|
||||
|
||||
# Execute
|
||||
result = self.context_caching.check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert - should return None after exhausting all pages without finding match
|
||||
assert result is None
|
||||
# Verify exactly MAX_PAGINATION_PAGES API calls were made (not more)
|
||||
assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]
|
||||
)
|
||||
@patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching")
|
||||
async def test_async_check_cache_pagination_max_pages_limit(
|
||||
self, mock_get_token_url, custom_llm_provider
|
||||
):
|
||||
"""Test that async pagination stops after MAX_PAGINATION_PAGES iterations"""
|
||||
# Setup
|
||||
mock_get_token_url.return_value = ("token", "https://test-url.com")
|
||||
cache_key_to_find = "nonexistent_cache_key"
|
||||
|
||||
# Create mock response that always has nextPageToken (infinite pagination scenario)
|
||||
def create_page_response(page_num):
|
||||
response = MagicMock()
|
||||
response.json.return_value = {
|
||||
"cachedContents": [
|
||||
{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}
|
||||
],
|
||||
"nextPageToken": f"token_page_{page_num + 1}",
|
||||
}
|
||||
return response
|
||||
|
||||
# Create MAX_PAGINATION_PAGES responses, each with a nextPageToken
|
||||
self.mock_async_client.get = AsyncMock(
|
||||
side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)]
|
||||
)
|
||||
|
||||
# Execute
|
||||
result = await self.context_caching.async_check_cache(
|
||||
cache_key=cache_key_to_find,
|
||||
client=self.mock_async_client,
|
||||
headers={"Authorization": "Bearer token"},
|
||||
api_key="test_key",
|
||||
api_base=None,
|
||||
logging_obj=self.mock_logging,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
vertex_project="test_project",
|
||||
vertex_location="us-central1",
|
||||
vertex_auth_header="Bearer test-token",
|
||||
)
|
||||
|
||||
# Assert - should return None after exhausting all pages without finding match
|
||||
assert result is None
|
||||
# Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more)
|
||||
assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES
|
||||
|
||||
|
||||
class TestVertexAIGlobalLocation:
|
||||
"""Test global location handling in context caching."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user