Merge commit 'abf9ff46e4' into litellm_azure_entra_token_followups

# Conflicts:
#	tests/unit/llms/azure/test_azure_common_utils.py

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:

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
@@ -2142,7 +2140,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(
@@ -2150,19 +2148,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",
@@ -2174,51 +2172,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_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_request(monkeypatch, api_version):
-    import threading
-
+async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_request(
+    monkeypatch: pytest.MonkeyPatch, api_version: str
+) -> None:
     import httpx
 
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -2279,7 +2261,9 @@ async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_re
 
 @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:

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
@@ -521,51 +520,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()