Skip to content

Commit 96aed6a

Browse files
hazyonehazyone
committed
test added
1 parent 311a0b3 commit 96aed6a

File tree

1 file changed

+218
-0
lines changed

1 file changed

+218
-0
lines changed

‎tests/test_litellm/proxy/pass_through_endpoints/test_anthropic_auth_headers.py‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
import os
2+
import sys
3+
from unittest.mock import AsyncMock, MagicMock, patch
4+
5+
import pytest
6+
7+
sys.path.insert(
8+
0, os.path.abspath("../../..")
9+
) # Adds the parent directory to the system path
10+
11+
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
12+
anthropic_proxy_route,
13+
)
14+
15+
16+
class TestAnthropicAuthHeaders:
17+
"""Test authentication header handling in anthropic_proxy_route."""
18+
19+
@pytest.fixture
20+
def mock_request(self):
21+
"""Create a mock request object."""
22+
request = MagicMock()
23+
request.method = "POST"
24+
request.headers = {}
25+
return request
26+
27+
@pytest.fixture
28+
def mock_response(self):
29+
"""Create a mock FastAPI response object."""
30+
return MagicMock()
31+
32+
@pytest.fixture
33+
def mock_user_api_key_dict(self):
34+
"""Create a mock user API key dict."""
35+
return {"user_id": "test_user"}
36+
37+
@pytest.mark.asyncio
38+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route")
39+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn")
40+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router")
41+
async def test_client_authorization_header_priority(
42+
self,
43+
mock_router,
44+
mock_streaming,
45+
mock_create_route,
46+
mock_request,
47+
mock_response,
48+
mock_user_api_key_dict,
49+
):
50+
"""Test that client Authorization header takes priority over server key."""
51+
# Setup
52+
mock_request.headers = {"authorization": "Bearer client-key-123"}
53+
mock_router.get_credentials.return_value = "server-key-456"
54+
mock_streaming.return_value = False
55+
mock_endpoint_func = AsyncMock(return_value="test_response")
56+
mock_create_route.return_value = mock_endpoint_func
57+
58+
# Act
59+
await anthropic_proxy_route(
60+
endpoint="v1/messages",
61+
request=mock_request,
62+
fastapi_response=mock_response,
63+
user_api_key_dict=mock_user_api_key_dict,
64+
)
65+
66+
# Assert
67+
mock_create_route.assert_called_once()
68+
call_kwargs = mock_create_route.call_args[1]
69+
70+
assert call_kwargs["custom_headers"] == {}
71+
assert call_kwargs["_forward_headers"] is True
72+
73+
@pytest.mark.asyncio
74+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route")
75+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn")
76+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router")
77+
async def test_client_x_api_key_header_priority(
78+
self,
79+
mock_router,
80+
mock_streaming,
81+
mock_create_route,
82+
mock_request,
83+
mock_response,
84+
mock_user_api_key_dict,
85+
):
86+
"""Test that client x-api-key header takes priority over server key."""
87+
# Setup
88+
mock_request.headers = {"x-api-key": "client-x-api-key-123"}
89+
mock_router.get_credentials.return_value = "server-key-456"
90+
mock_streaming.return_value = False
91+
mock_endpoint_func = AsyncMock(return_value="test_response")
92+
mock_create_route.return_value = mock_endpoint_func
93+
94+
# Act
95+
await anthropic_proxy_route(
96+
endpoint="v1/messages",
97+
request=mock_request,
98+
fastapi_response=mock_response,
99+
user_api_key_dict=mock_user_api_key_dict,
100+
)
101+
102+
# Assert
103+
mock_create_route.assert_called_once()
104+
call_kwargs = mock_create_route.call_args[1]
105+
106+
assert call_kwargs["custom_headers"] == {}
107+
assert call_kwargs["_forward_headers"] is True
108+
109+
@pytest.mark.asyncio
110+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route")
111+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn")
112+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router")
113+
async def test_server_api_key_fallback(
114+
self,
115+
mock_router,
116+
mock_streaming,
117+
mock_create_route,
118+
mock_request,
119+
mock_response,
120+
mock_user_api_key_dict,
121+
):
122+
"""Test that server API key is used when no client authentication is provided."""
123+
# Setup
124+
mock_request.headers = {} # No authentication headers
125+
mock_router.get_credentials.return_value = "server-key-456"
126+
mock_streaming.return_value = False
127+
mock_endpoint_func = AsyncMock(return_value="test_response")
128+
mock_create_route.return_value = mock_endpoint_func
129+
130+
# Act
131+
await anthropic_proxy_route(
132+
endpoint="v1/messages",
133+
request=mock_request,
134+
fastapi_response=mock_response,
135+
user_api_key_dict=mock_user_api_key_dict,
136+
)
137+
138+
# Assert
139+
mock_create_route.assert_called_once()
140+
call_kwargs = mock_create_route.call_args[1]
141+
142+
assert call_kwargs["custom_headers"] == {"x-api-key": "server-key-456"}
143+
assert call_kwargs["_forward_headers"] is True
144+
145+
@pytest.mark.asyncio
146+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route")
147+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn")
148+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router")
149+
async def test_no_authentication_available(
150+
self,
151+
mock_router,
152+
mock_streaming,
153+
mock_create_route,
154+
mock_request,
155+
mock_response,
156+
mock_user_api_key_dict,
157+
):
158+
"""Test that no x-api-key header is added when no authentication is available."""
159+
# Setup
160+
mock_request.headers = {} # No authentication headers
161+
mock_router.get_credentials.return_value = None # No server key
162+
mock_streaming.return_value = False
163+
mock_endpoint_func = AsyncMock(return_value="test_response")
164+
mock_create_route.return_value = mock_endpoint_func
165+
166+
# Act
167+
await anthropic_proxy_route(
168+
endpoint="v1/messages",
169+
request=mock_request,
170+
fastapi_response=mock_response,
171+
user_api_key_dict=mock_user_api_key_dict,
172+
)
173+
174+
# Assert
175+
mock_create_route.assert_called_once()
176+
call_kwargs = mock_create_route.call_args[1]
177+
178+
assert call_kwargs["custom_headers"] == {}
179+
assert call_kwargs["_forward_headers"] is True
180+
181+
@pytest.mark.asyncio
182+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route")
183+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn")
184+
@patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router")
185+
async def test_both_client_headers_present(
186+
self,
187+
mock_router,
188+
mock_streaming,
189+
mock_create_route,
190+
mock_request,
191+
mock_response,
192+
mock_user_api_key_dict,
193+
):
194+
"""Test that no server key is added when client has both auth headers."""
195+
# Setup
196+
mock_request.headers = {
197+
"authorization": "Bearer client-auth-key",
198+
"x-api-key": "client-x-api-key"
199+
}
200+
mock_router.get_credentials.return_value = "server-key-456"
201+
mock_streaming.return_value = False
202+
mock_endpoint_func = AsyncMock(return_value="test_response")
203+
mock_create_route.return_value = mock_endpoint_func
204+
205+
# Act
206+
await anthropic_proxy_route(
207+
endpoint="v1/messages",
208+
request=mock_request,
209+
fastapi_response=mock_response,
210+
user_api_key_dict=mock_user_api_key_dict,
211+
)
212+
213+
# Assert
214+
mock_create_route.assert_called_once()
215+
call_kwargs = mock_create_route.call_args[1]
216+
217+
assert call_kwargs["custom_headers"] == {}
218+
assert call_kwargs["_forward_headers"] is True

0 commit comments

Comments
 (0)