test(test_anthropic_experimental_pass_through_messages_handler.py): add unit test confirming custom llm respected

This commit is contained in:
Krrish Dholakia 2025-06-24 14:05:42 -07:00
parent 1ecf0315e5
commit ded1734d34

View file

@ -34,3 +34,33 @@ def test_anthropic_experimental_pass_through_messages_handler():
print(f"Error: {e}")
mock_completion.assert_called_once()
mock_completion.call_args.kwargs["api_key"] == "test-api-key"
def test_anthropic_experimental_pass_through_messages_handler_custom_llm_provider():
"""
Test that litellm.completion is called when a custom LLM provider is given
"""
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages_handler,
)
with patch("litellm.completion", return_value="test-response") as mock_completion:
try:
anthropic_messages_handler(
max_tokens=100,
messages=[{"role": "user", "content": "Hello, how are you?"}],
model="my-custom-model",
custom_llm_provider="my-custom-llm",
api_key="test-api-key",
)
except Exception as e:
print(f"Error: {e}")
# Assert that litellm.completion was called when using a custom LLM provider
mock_completion.assert_called_once()
# Verify that the custom provider was passed through
call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["custom_llm_provider"] == "my-custom-llm"
assert call_kwargs["model"] == "my-custom-llm/my-custom-model"
assert call_kwargs["api_key"] == "test-api-key"