mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
test(anthropic): add unit test coverage for count_tokens custom api_base endpoint resolution
This commit is contained in:
parent
955f6d8cae
commit
0569cbe8d8
1 changed files with 32 additions and 0 deletions
32
tests/test_anthropic_count_tokens.py
Normal file
32
tests/test_anthropic_count_tokens.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_count_tokens_custom_api_base():
|
||||
handler = AnthropicCountTokensHandler()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"input_tokens": 10}
|
||||
|
||||
with patch("litellm.llms.anthropic.count_tokens.handler.get_async_httpx_client") as mock_client_factory:
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = mock_response
|
||||
mock_client_factory.return_value = mock_client
|
||||
|
||||
# Test custom api_base URL resolution and formatting
|
||||
res = await handler.handle_count_tokens_request(
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
api_key="test-key",
|
||||
api_base="https://custom-proxy.internal.com/v1",
|
||||
)
|
||||
|
||||
assert res["input_tokens"] == 10
|
||||
mock_client.post.assert_called_once()
|
||||
called_url = mock_client.post.call_args[0][0]
|
||||
assert called_url == "https://custom-proxy.internal.com/v1/messages/count_tokens"
|
||||
Loading…
Add table
Reference in a new issue