From d231d0f5b360f951106e7fba404a73ff1dfb0ce6 Mon Sep 17 00:00:00 2001 From: Gabi Davar Date: Fri, 16 Oct 2026 23:04:19 +0300 Subject: [PATCH] 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 --- litellm/llms/vertex_ai/vertex_llm_base.py | 7 +- .../llms/vertex_ai/test_vertex_llm_base.py | 136 +++++++++++------- 2 files changed, 89 insertions(+), 54 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 085bb55dc59..c1dccefca37 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -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 diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index 1781d3d8226..2baec4ea74e 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -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(