mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-24 05:45:11 +00:00
test(vertex_ai): drive the expiry margin tests from a frozen clock and Google's loader
The two regression tests patched VertexBase.load_auth and read the real clock, and the background refresh assertion raced a 50ms sleep. They now patch google.oauth2's from_service_account_info with a credential whose token state follows one frozen clock, the module reads that clock through time.time so the test freezes it without touching LiteLLM code, and the background refresh is awaited from the tracked task instead of slept on
This commit is contained in:
Gabi Davar
parent
00b4940899
commit
d231d0f5b3
2 changed files with 89 additions and 54 deletions
|
|
@ -7,6 +7,7 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
|
|
@ -71,6 +72,10 @@ class _VertexCredentialsObject(Protocol):
|
|||
def refresh(self, request: object) -> None: ...
|
||||
|
||||
|
||||
def _utcnow() -> datetime:
|
||||
return datetime.fromtimestamp(time.time(), tz=timezone.utc)
|
||||
|
||||
|
||||
class VertexBase:
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
|
@ -499,7 +504,7 @@ class VertexBase:
|
|||
expiry: Final = getattr(credentials, "expiry", None)
|
||||
if not isinstance(expiry, datetime):
|
||||
return False
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
now: Final = _utcnow()
|
||||
remaining: Final = expiry - (now if expiry.tzinfo is not None else now.replace(tzinfo=None))
|
||||
return remaining.total_seconds() <= VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
|
@ -2201,6 +2202,8 @@ class TestVertexBase:
|
|||
|
||||
assert token == "cached-token"
|
||||
assert not mock_get_lock.called, "Fast path should not acquire lock"
|
||||
assert not mock_get_lock.called, "Fast path should not acquire lock"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_token_inside_safety_margin_refreshes_before_use(self):
|
||||
|
|
@ -2209,45 +2212,30 @@ class TestVertexBase:
|
|||
a background refresh fails the request that triggered the refresh."""
|
||||
import datetime
|
||||
|
||||
from google.auth import _helpers as google_auth_helpers
|
||||
from google.auth.credentials import TokenState
|
||||
|
||||
from litellm.constants import VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS
|
||||
|
||||
clock = _FakeClock(datetime.datetime(2026, 1, 1, 12, 0, 0))
|
||||
creds = _FakeGoogleCredentials(clock)
|
||||
vertex_base = VertexBase()
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.token = "about-to-expire-token"
|
||||
mock_creds.token_state = TokenState.STALE
|
||||
mock_creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta(
|
||||
seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS / 2
|
||||
)
|
||||
mock_creds.project_id = "project-1"
|
||||
mock_creds.quota_project_id = "project-1"
|
||||
|
||||
credentials = {"type": "service_account", "project_id": "project-1"}
|
||||
|
||||
with (
|
||||
patch.object(vertex_base, "load_auth", return_value=(mock_creds, "project-1")),
|
||||
patch.object(vertex_base, "refresh_auth") as mock_refresh,
|
||||
patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds),
|
||||
patch("time.time", side_effect=clock.timestamp),
|
||||
):
|
||||
|
||||
def mock_refresh_impl(creds):
|
||||
creds.token = "refreshed-token"
|
||||
creds.token_state = TokenState.FRESH
|
||||
creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta(hours=1)
|
||||
|
||||
mock_refresh.side_effect = mock_refresh_impl
|
||||
first_token, _ = await vertex_base.ensure_access_token_async(
|
||||
credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
clock.now = creds.expiry - datetime.timedelta(seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS / 2)
|
||||
|
||||
token, project = await vertex_base.ensure_access_token_async(
|
||||
credentials=credentials,
|
||||
project_id="project-1",
|
||||
custom_llm_provider="vertex_ai",
|
||||
credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
|
||||
assert token == "refreshed-token"
|
||||
assert project == "project-1"
|
||||
assert mock_refresh.call_count == 1
|
||||
assert first_token == "token-1"
|
||||
assert token == "token-2"
|
||||
assert project == "project-1"
|
||||
assert creds.refresh_calls == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_token_outside_safety_margin_is_served_while_refreshing(self):
|
||||
|
|
@ -2255,45 +2243,87 @@ class TestVertexBase:
|
|||
import datetime
|
||||
|
||||
from google.auth import _helpers as google_auth_helpers
|
||||
from google.auth.credentials import TokenState
|
||||
|
||||
from litellm.constants import VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS
|
||||
|
||||
stale_window = google_auth_helpers.REFRESH_THRESHOLD.total_seconds()
|
||||
if VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS >= stale_window:
|
||||
pytest.skip("the configured margin covers the whole STALE window, there is no zero-latency path to test")
|
||||
seconds_left = (VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS + stale_window) / 2
|
||||
|
||||
clock = _FakeClock(datetime.datetime(2026, 1, 1, 12, 0, 0))
|
||||
creds = _FakeGoogleCredentials(clock)
|
||||
vertex_base = VertexBase()
|
||||
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.token = "near-expiry-token"
|
||||
mock_creds.token_state = TokenState.STALE
|
||||
mock_creds.expiry = google_auth_helpers.utcnow() + datetime.timedelta(
|
||||
seconds=VERTEX_AI_TOKEN_EXPIRY_SAFETY_MARGIN_SECONDS * 2
|
||||
)
|
||||
mock_creds.project_id = "project-1"
|
||||
mock_creds.quota_project_id = "project-1"
|
||||
|
||||
credentials = {"type": "service_account", "project_id": "project-1"}
|
||||
|
||||
with (
|
||||
patch.object(vertex_base, "load_auth", return_value=(mock_creds, "project-1")),
|
||||
patch.object(vertex_base, "refresh_auth") as mock_refresh,
|
||||
patch("google.oauth2.service_account.Credentials.from_service_account_info", return_value=creds),
|
||||
patch("time.time", side_effect=clock.timestamp),
|
||||
):
|
||||
|
||||
def mock_refresh_impl(creds):
|
||||
creds.token = "refreshed-token"
|
||||
creds.token_state = TokenState.FRESH
|
||||
|
||||
mock_refresh.side_effect = mock_refresh_impl
|
||||
await vertex_base.ensure_access_token_async(
|
||||
credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
clock.now = creds.expiry - datetime.timedelta(seconds=seconds_left)
|
||||
|
||||
token, _ = await vertex_base.ensure_access_token_async(
|
||||
credentials=credentials,
|
||||
project_id="project-1",
|
||||
custom_llm_provider="vertex_ai",
|
||||
credentials=credentials, project_id="project-1", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
assert token == "token-1"
|
||||
await asyncio.gather(*vertex_base._background_refresh_tasks.values())
|
||||
|
||||
assert token == "near-expiry-token"
|
||||
assert creds.refresh_calls == 2
|
||||
assert creds.token == "token-2"
|
||||
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert mock_refresh.called
|
||||
class _FakeClock:
|
||||
"""The one clock both the module (through time.time) and the fake credential read."""
|
||||
|
||||
def __init__(self, now: datetime.datetime):
|
||||
self.now = now
|
||||
|
||||
def timestamp(self) -> float:
|
||||
return self.now.replace(tzinfo=datetime.timezone.utc).timestamp()
|
||||
|
||||
|
||||
class _FakeGoogleCredentials:
|
||||
"""Stands in for google.oauth2 credentials: the frozen clock decides FRESH / STALE / INVALID
|
||||
the way google-auth does, and refresh() mints the next token an hour out."""
|
||||
|
||||
def __init__(self, clock: _FakeClock):
|
||||
import datetime
|
||||
|
||||
self._clock = clock
|
||||
self._refresh_threshold = datetime.timedelta(minutes=3, seconds=45)
|
||||
self.token: str | None = None
|
||||
self.expiry = clock.now
|
||||
self.refresh_calls = 0
|
||||
self.project_id = "project-1"
|
||||
self.quota_project_id = "project-1"
|
||||
|
||||
@property
|
||||
def token_state(self):
|
||||
from google.auth.credentials import TokenState
|
||||
|
||||
if self.token is None or self._clock.now >= self.expiry:
|
||||
return TokenState.INVALID
|
||||
if self._clock.now >= self.expiry - self._refresh_threshold:
|
||||
return TokenState.STALE
|
||||
return TokenState.FRESH
|
||||
|
||||
@property
|
||||
def expired(self) -> bool:
|
||||
return self._clock.now >= self.expiry - self._refresh_threshold
|
||||
|
||||
@property
|
||||
def valid(self) -> bool:
|
||||
return self.token is not None and not self.expired
|
||||
|
||||
def refresh(self, request) -> None:
|
||||
import datetime
|
||||
|
||||
self.refresh_calls += 1
|
||||
self.token = f"token-{self.refresh_calls}"
|
||||
self.expiry = self._clock.now + datetime.timedelta(hours=1)
|
||||
|
||||
|
||||
_ENCODED_VERTEX_KEY: Final = base64.b64encode(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue