mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
"""Tests for litellm/llms/a2a/common_utils.py."""
|
|
|
|
from collections.abc import Mapping
|
|
from types import MappingProxyType
|
|
|
|
import pytest
|
|
|
|
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
|
|
|
|
|
|
class _RecordingEntraResolver:
|
|
def __init__(self) -> None:
|
|
self.calls: list[Mapping[str, object]] = []
|
|
|
|
async def __call__(self, litellm_params: Mapping[str, object]) -> Mapping[str, str]:
|
|
self.calls.append(litellm_params)
|
|
return MappingProxyType({"Authorization": "Bearer minted-entra-token"})
|
|
|
|
|
|
_SERVICE_PRINCIPAL = MappingProxyType({"tenant_id": "tenant", "client_id": "client", "client_secret": "sp-secret"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entra_agent_gets_a_minted_bearer_for_the_a2a_hop():
|
|
resolver = _RecordingEntraResolver()
|
|
|
|
header = await resolve_a2a_hop_auth_header(_SERVICE_PRINCIPAL, None, resolver)
|
|
|
|
assert header == {"Authorization": "Bearer minted-entra-token"}
|
|
assert resolver.calls == [_SERVICE_PRINCIPAL]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completion_bridge_agent_keeps_its_entra_credentials_for_the_model_provider():
|
|
"""A bridged agent's tenant_id/client_id/client_secret authenticate the model it bridges to, so the A2A hop
|
|
must not spend them on a bearer of its own."""
|
|
resolver = _RecordingEntraResolver()
|
|
|
|
header = await resolve_a2a_hop_auth_header(_SERVICE_PRINCIPAL, "azure_ai", resolver)
|
|
|
|
assert header is None
|
|
assert resolver.calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_without_entra_credentials_gets_no_bearer():
|
|
resolver = _RecordingEntraResolver()
|
|
|
|
header = await resolve_a2a_hop_auth_header({"api_base": "https://agent.example.com"}, None, resolver)
|
|
|
|
assert header is None
|
|
assert resolver.calls == []
|