From abf9ff46e46073a04c64ab851c8e4b03244f5147 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:59:54 -0700 Subject: [PATCH] test(azure): assert the token and validation thread hops by thread identity instead of sleeps --- ...st_bing_grounding_search_transformation.py | 11 ++-- .../llms/azure/test_azure_common_utils.py | 52 +++++++------------ .../custom_httpx/test_llm_http_handler.py | 50 ++++++------------ 3 files changed, 41 insertions(+), 72 deletions(-) diff --git a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py index 5a97b6f933c..dfd1474211f 100644 --- a/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/unit/llms/azure/search/test_bing_grounding_search_transformation.py @@ -1,5 +1,6 @@ import json, litellm from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, Mock, patch import pytest @@ -110,21 +111,21 @@ def test_validate_environment_entra_minter_failure_names_the_options(): assert "no az login" in message -def test_default_entra_minter_reuses_provider(monkeypatch: pytest.MonkeyPatch): +def test_default_entra_minter_reuses_provider(monkeypatch: pytest.MonkeyPatch) -> None: import importlib from litellm.llms.azure import common_utils from litellm.llms.azure.search.transformation import _default_entra_token_minter - provider_module = importlib.import_module("litellm.secret_managers.get_azure_ad_token_provider") - provider_factory = Mock(side_effect=lambda **_: lambda: "entra-token") + provider_module: Final = importlib.import_module("litellm.secret_managers.get_azure_ad_token_provider") + provider_factory: Final = Mock(side_effect=lambda **_: lambda: "entra-token") common_utils._cached_azure_ad_token_refresh_provider.cache_clear() monkeypatch.setattr(common_utils, "get_azure_ad_token_provider", provider_factory) monkeypatch.setattr(provider_module, "get_azure_ad_token_provider", provider_factory) try: - first = _default_entra_token_minter() - second = _default_entra_token_minter() + first: Final = _default_entra_token_minter() + second: Final = _default_entra_token_minter() assert first == second == "entra-token" assert provider_factory.call_count == 1 finally: diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index c60e3e8a9d8..a09e4c72bd8 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1,8 +1,6 @@ -import asyncio -import itertools import json import os -import time +import threading import traceback from typing import Callable, Final, Optional from unittest.mock import MagicMock, patch @@ -2145,7 +2143,7 @@ class TestEntraIdTokenProviderCache: assert mock_credential.call_count == 2 -def test_username_password_token_provider_reuses_same_credentials_and_separates_passwords(): +def test_username_password_token_provider_reuses_same_credentials_and_separates_passwords() -> None: with ( patch("azure.identity.UsernamePasswordCredential") as mock_credential, patch( @@ -2153,19 +2151,19 @@ def test_username_password_token_provider_reuses_same_credentials_and_separates_ side_effect=lambda credential, scope: lambda: "token", ), ): - first = get_azure_ad_token_from_username_password( + first: Final = get_azure_ad_token_from_username_password( client_id="client", azure_username="user", azure_password="password", scope="https://cognitiveservices.azure.com/.default", ) - second = get_azure_ad_token_from_username_password( + second: Final = get_azure_ad_token_from_username_password( client_id="client", azure_username="user", azure_password="password", scope="https://cognitiveservices.azure.com/.default", ) - changed_password = get_azure_ad_token_from_username_password( + changed_password: Final = get_azure_ad_token_from_username_password( client_id="client", azure_username="user", azure_password="different-password", @@ -2177,49 +2175,35 @@ def test_username_password_token_provider_reuses_same_credentials_and_separates_ @pytest.mark.asyncio -async def test_non_v1_azure_token_refresh_does_not_block_the_event_loop(): +async def test_non_v1_azure_token_refresh_runs_off_the_event_loop_thread() -> None: from openai import AsyncAzureOpenAI - base_llm = BaseAzureLLM() - api_base: Final = "https://test.openai.azure.com" + def token_provider_naming_its_thread() -> str: + return str(threading.get_ident()) - def slow_token_provider() -> str: - time.sleep(0.3) - return "slow-token" - - client: Final = base_llm.get_azure_openai_client( + client: Final = BaseAzureLLM().get_azure_openai_client( api_key=None, - api_base=api_base, + api_base="https://test.openai.azure.com", api_version="2024-02-01", - litellm_params={"azure_ad_token_provider": slow_token_provider}, + litellm_params={"azure_ad_token_provider": token_provider_naming_its_thread}, _is_async=True, ) assert isinstance(client, AsyncAzureOpenAI) - ticks: Final = itertools.count() - token_resolved: Final = asyncio.Event() - - async def resolve_token() -> None: - try: - assert await client._get_azure_ad_token() == "slow-token" - finally: - token_resolved.set() - - async def ticker() -> None: - while not token_resolved.is_set(): - next(ticks) - await asyncio.sleep(0.01) - try: - await asyncio.gather(resolve_token(), ticker()) - assert next(ticks) > 0 + token_thread: Final = await client._get_azure_ad_token() finally: await client.close() + assert token_thread is not None + assert token_thread != str(threading.get_ident()) + @pytest.mark.asyncio @pytest.mark.parametrize("api_version", ["2024-02-01", "v1"]) -async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token(monkeypatch, api_version): +async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token( + monkeypatch: pytest.MonkeyPatch, api_version: str +) -> None: import httpx def echo_authorization(request: httpx.Request) -> httpx.Response: diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 0d4825152cb..6504d06fa92 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1,7 +1,6 @@ import asyncio import base64 import inspect -import itertools import json import logging import threading @@ -478,51 +477,36 @@ async def test_async_response_api_handler_streams_when_provider_transform_adds_s @pytest.mark.asyncio -async def test_async_response_api_handler_validates_environment_off_the_event_loop(): - validation_started: Final = threading.Event() - validation_finished: Final = threading.Event() - ticks: Final = itertools.count() - +async def test_async_response_api_handler_validates_environment_off_the_event_loop() -> None: def validate_environment(**_kwargs: object) -> dict[str, str]: - validation_started.set() - time.sleep(0.3) - validation_finished.set() - return {} + return {"validated-on-thread": str(threading.get_ident())} - config = Mock() + config: Final = Mock() config.validate_environment.side_effect = validate_environment config.get_complete_url.return_value = "https://responses.example.com/v1/responses" config.async_transform_responses_api_request = AsyncMock(return_value={"model": "model", "input": "hi"}) config.sign_request.return_value = ({}, None) - client = AsyncMock(spec=AsyncHTTPHandler) + client: Final = AsyncMock(spec=AsyncHTTPHandler) client.post.return_value = httpx.Response( 200, request=httpx.Request("POST", "https://responses.example.com/v1/responses"), ) - logging_obj = Mock() + logging_obj: Final = Mock() logging_obj.dynamic_success_callbacks = None - async def call_handler() -> None: - await BaseLLMHTTPHandler().async_response_api_handler( - model="model", - input="hi", - responses_api_provider_config=config, - response_api_optional_request_params={}, - custom_llm_provider="test", - litellm_params=GenericLiteLLMParams(), - logging_obj=logging_obj, - client=client, - ) + await BaseLLMHTTPHandler().async_response_api_handler( + model="model", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="test", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) - async def ticker() -> None: - while not validation_finished.is_set(): - if validation_started.is_set(): - next(ticks) - await asyncio.sleep(0.01) - - await asyncio.gather(call_handler(), ticker()) - - assert next(ticks) > 0 + validated_headers: Final = config.async_transform_responses_api_request.await_args.kwargs["headers"] + assert validated_headers["validated-on-thread"] != str(threading.get_ident()) client.post.assert_awaited_once()