mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
e1d2789d29
Migrate 18 provider test files from tests/test_litellm/llms to tests/unit/llms. 194 kept tests move as-is after mutation testing; 1 test deleted (test_completion_datarobot_with_environment_variables, env-gated no-assert); the fixture-only fal_ai cost calculator file is removed. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
import json
|
|
import os
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from litellm import completion
|
|
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
|
|
|
|
|
@patch.dict(os.environ, {}, clear=True)
|
|
def test_completion_datarobot():
|
|
"""Ensure that the completion function works with DataRobot API."""
|
|
messages = [
|
|
{"role": "user", "content": "What's the weather like in San Francisco?"}
|
|
]
|
|
try:
|
|
client = HTTPHandler()
|
|
with patch.object(client, "post") as mock_post:
|
|
response = completion(
|
|
model="datarobot/vertex_ai/gemini-1.5-flash-002",
|
|
messages=messages,
|
|
client=client,
|
|
max_tokens=5,
|
|
clientId="custom-model",
|
|
)
|
|
print(response)
|
|
|
|
# Add any assertions here to check the response
|
|
mock_post.assert_called_once()
|
|
mocks_kwargs = mock_post.call_args.kwargs
|
|
assert (
|
|
mocks_kwargs["url"]
|
|
== "https://app.datarobot.com/api/v2/genai/llmgw/chat/completions/"
|
|
)
|
|
assert mocks_kwargs["headers"]["Authorization"] == "Bearer fake-api-key"
|
|
json_data = json.loads(mock_post.call_args.kwargs["data"])
|
|
assert json_data["clientId"] == "custom-model"
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@patch.dict(
|
|
os.environ,
|
|
{
|
|
"DATAROBOT_ENDPOINT": "https://app.datarobot.com/api/v2/deployments/deployment_id/"
|
|
},
|
|
clear=True,
|
|
)
|
|
def test_completion_datarobot_with_deployment():
|
|
"""Ensure that deployment URL is used correctly."""
|
|
messages = [
|
|
{"role": "user", "content": "What's the weather like in San Francisco?"}
|
|
]
|
|
try:
|
|
client = HTTPHandler()
|
|
with patch.object(client, "post") as mock_post:
|
|
response = completion(
|
|
model="datarobot/vertex_ai/gemini-1.5-flash-002",
|
|
messages=messages,
|
|
client=client,
|
|
max_tokens=5,
|
|
clientId="custom-model",
|
|
)
|
|
print(response)
|
|
|
|
# Add any assertions here to check the response
|
|
mock_post.assert_called_once()
|
|
mocks_kwargs = mock_post.call_args.kwargs
|
|
assert (
|
|
mocks_kwargs["url"]
|
|
== "https://app.datarobot.com/api/v2/deployments/deployment_id/"
|
|
)
|
|
assert mocks_kwargs["headers"]["Authorization"] == "Bearer fake-api-key"
|
|
json_data = json.loads(mock_post.call_args.kwargs["data"])
|
|
assert json_data["clientId"] == "custom-model"
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|