scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Merge commit '6bcd5b827a' into litellm_azure_entra_token_followups

Browse source
This commit is contained in:
mateo-berri 2026-10-19 12:11:10 -07:00
commit 0eee9e1ba0
1 changed files with 21 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
tests/unit/llms/azure/test_azure_common_utils.py
View file

@ -2,11 +2,11 @@ import json
import os
import threading
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Callable, Final, Optional
from urllib.parse import parse_qs
from unittest.mock import MagicMock, patch
import httpx
import pytest
import litellm
@ -2287,8 +2287,6 @@ async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_re
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:
return httpx.Response(
200,
@ -2326,44 +2324,32 @@ async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token(
assert response.choices[0].message.content == "Bearer async-token"
def test_oidc_token_exchange_closes_its_token_endpoint_connection(monkeypatch: pytest.MonkeyPatch) -> None:
closed_connections: Final = threading.Semaphore(0)
def test_oidc_token_exchange_closes_its_own_http_client(monkeypatch: pytest.MonkeyPatch) -> None:
class TokenEndpoint(httpx.BaseTransport):
def __init__(self) -> None:
self.closed = False
class TokenEndpoint(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def handle_request(self, request: httpx.Request) -> httpx.Response:
form: Final = parse_qs(request.read().decode())
return httpx.Response(
200, json={"access_token": f"minted-for-{form['client_assertion'][0]}", "expires_in": 3600}
)
def do_POST(self) -> None:
form: Final = parse_qs(self.rfile.read(int(self.headers["Content-Length"])).decode())
body: Final = json.dumps(
{"access_token": f"minted-for-{form['client_assertion'][0]}", "expires_in": 3600}
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def close(self) -> None:
self.closed = True
def finish(self) -> None:
super().finish()
closed_connections.release()
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), TokenEndpoint)
threading.Thread(target=server.serve_forever, daemon=True).start()
monkeypatch.setenv("AZURE_AUTHORITY_HOST", f"http://127.0.0.1:{server.server_port}")
token_endpoint: Final = TokenEndpoint()
monkeypatch.setattr(litellm, "sync_transport", token_endpoint, raising=False)
monkeypatch.setenv("AZURE_TEST_OIDC_ASSERTION", "workload-jwt")
try:
token: Final = get_azure_ad_token_from_oidc(
azure_ad_token="os.environ/AZURE_TEST_OIDC_ASSERTION",
azure_client_id="client",
azure_tenant_id="tenant",
)
connection_closed: Final = closed_connections.acquire(timeout=5)
finally:
server.shutdown()
server.server_close()
token: Final = get_azure_ad_token_from_oidc(
azure_ad_token="os.environ/AZURE_TEST_OIDC_ASSERTION",
azure_client_id="client",
azure_tenant_id="tenant-closes-its-client",
)
assert token == "minted-for-workload-jwt"
assert connection_closed
assert token_endpoint.closed
def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch):