scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

test: migrate databricks, xai, cometapi and fireworks translation tests

Browse source
This commit is contained in:
yuneng 2026-10-10 05:43:16 +00:00
parent 443797e7f5
commit 5d7d02fee5
5 changed files with 214 additions and 233 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

113
tests/llm_translation/test_cometapi_chat_transformation.py
View file

@ -1,113 +0,0 @@
"""
Unit tests for CometAPI Chat Configuration
Tests the CometAPIChatConfig class methods using mocks
"""
import os
import pytest
# Integration test example (requires real API key)
def test_cometapi_streaming_integration():
"""
Integration test for streaming - requires real API key
Run with: pytest -k test_cometapi_streaming_integration -s
"""
from litellm import completion
# Try to get API key from multiple environment variables
api_key = (
os.getenv("COMETAPI_API_KEY")
or os.getenv("COMETAPI_KEY")
or os.getenv("COMET_API_KEY")
)
if not api_key:
pytest.skip("COMETAPI_API_KEY not set - skipping streaming integration test")
try:
print(
f"🔍 Testing streaming with API key: {api_key[:6]}...{api_key[-4:]} (length: {len(api_key)})"
)
print(f"🔍 API base URL: {os.getenv('COMETAPI_API_BASE', 'default')}")
# test streaming API call
response = completion(
model="cometapi/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Count from 1 to 5"}],
api_key=api_key,
max_tokens=50,
stream=True,
)
# collect streaming response
chunks = []
content_parts = []
for chunk in response:
chunks.append(chunk)
if chunk.choices[0].delta.content:
content_parts.append(chunk.choices[0].delta.content)
# Verify we received at least one chunk and content
assert len(chunks) > 0, "Should receive at least one chunk"
assert len(content_parts) > 0, "Should receive content in chunks"
full_content = "".join(content_parts)
assert len(full_content.strip()) > 0, "Should have non-empty content"
print(f"✅ Received {len(chunks)} chunks")
print(f"✅ Full content: {full_content}")
except Exception as e:
print(f"❌ Streaming integration test error details:")
print(f" Error type: {type(e).__name__}")
print(f" Error message: {str(e)}")
if hasattr(e, "status_code"):
print(f" Status code: {e.status_code}")
if hasattr(e, "response"):
print(f" Response: {e.response}")
# Re-raise with more context for pytest
pytest.fail(f"Streaming integration test failed: {type(e).__name__}: {str(e)}")
def test_cometapi_with_custom_base_url():
"""
Test CometAPI with custom base URL
"""
from litellm import completion
api_key = (
os.getenv("COMETAPI_API_KEY")
or os.getenv("COMETAPI_KEY")
or os.getenv("COMET_API_KEY")
)
custom_base_url = os.getenv("COMETAPI_API_BASE", "https://api.cometapi.com/v1")
if not api_key:
pytest.skip("COMETAPI_API_KEY not set - skipping custom base URL test")
try:
response = completion(
model="cometapi/gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello"}],
api_key=api_key,
api_base=custom_base_url,
max_tokens=5,
)
assert response.choices[0].message.content
print(f"✅ Custom base URL test passed: {response.choices[0].message.content}")
except Exception as e:
pytest.fail(f"Custom base URL test failed: {str(e)}")
if __name__ == "__main__":
# Quick test runner
pytest.main([__file__, "-v"])

118
tests/llm_translation/test_xai.py
View file

@ -1,118 +0,0 @@
import json
import os
from datetime import datetime
from unittest.mock import AsyncMock, patch
import httpx
import pytest
from base_llm_unit_tests import BaseLLMChatTest, BaseReasoningLLMTests
from litellm import Choices, EmbeddingResponse, Message, ModelResponse, Usage, completion
from litellm.llms.xai.chat.transformation import XAI_API_BASE, XAIChatConfig
def test_xai_message_name_filtering():
messages = [
{
"role": "system",
"content": "*I press the green button*",
"name": "example_user",
},
{"role": "user", "content": "Hello", "name": "John"},
{"role": "assistant", "content": "Hello", "name": "Jane"},
]
response = completion(
model="xai/grok-3-mini-beta",
messages=messages,
)
assert response is not None
assert response.choices[0].message.content is not None
class TestXAIReasoningEffort(BaseReasoningLLMTests):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-4.6",
"messages": [
{"role": "user", "content": "What is 17 times 23? Work it out step by step, then give the number."}
],
}
class TestXAIChat(BaseLLMChatTest):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-3-mini-beta",
}
test_web_search = None
def test_xai_streaming_with_include_usage():
"""
Test that xAI streaming correctly handles usage in the last chunk
when stream_options={"include_usage": True} is set.
xAI sends usage in a chunk with empty choices array, which should be
handled by XAIChatCompletionStreamingHandler.
"""
try:
response = completion(
model="xai/grok-4-1-fast-non-reasoning",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Say hello in one word"},
],
stream=True,
stream_options={"include_usage": True},
max_tokens=10,
)
chunks = []
usage_chunk = None
for chunk in response:
chunks.append(chunk)
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk
# Verify we got chunks
assert len(chunks) > 0, "Should receive streaming chunks"
# Verify usage was included in one of the chunks
assert usage_chunk is not None, "Should receive usage in streaming chunks"
# Verify usage has expected fields
assert hasattr(
usage_chunk.usage, "prompt_tokens"
), "Usage should have prompt_tokens"
assert hasattr(
usage_chunk.usage, "completion_tokens"
), "Usage should have completion_tokens"
assert hasattr(
usage_chunk.usage, "total_tokens"
), "Usage should have total_tokens"
# Verify usage values are positive
assert usage_chunk.usage.prompt_tokens > 0, "prompt_tokens should be positive"
assert (
usage_chunk.usage.completion_tokens > 0
), "completion_tokens should be positive"
assert usage_chunk.usage.total_tokens > 0, "total_tokens should be positive"
print(f"✓ Successfully received usage in streaming chunk: {usage_chunk.usage}")
except Exception as e:
if "API key" in str(e) or "authentication" in str(e).lower():
pytest.skip(f"Skipping test due to API key issue: {str(e)}")
raise

43
tests/unit/llms/cometapi/chat/test_cometapi_chat_transformation.py
View file

@ -4,10 +4,14 @@ Unit tests for CometAPI Chat Configuration
Tests the CometAPIChatConfig class methods using mocks
"""
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
from litellm.llms.cometapi.chat.transformation import (
CometAPIChatCompletionStreamingHandler,
CometAPIConfig,
@ -44,6 +48,7 @@ class TestCometAPIChatCompletionStreamingHandler:
assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"]
assert result.usage.total_tokens == chunk["usage"]["total_tokens"]
assert len(result.choices) == 1
assert result.choices[0]["delta"]["content"] == "test content"
assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning"
def test_chunk_parser_error_response(self):
@ -175,6 +180,42 @@ class TestCometAPIConfig:
assert error.status_code == 400
@pytest.mark.respx(assert_all_called=True)
def test_completion_uses_custom_api_base(respx_mock: respx.MockRouter) -> None:
api_base: Final = "https://cometapi.example/custom/v1"
route: Final = respx_mock.post(f"{api_base}/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-cometapi",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-5.6",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "OK"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
},
)
)
response: Final = litellm.completion(
model="cometapi/gpt-5.6",
messages=[{"role": "user", "content": "Say OK"}],
api_key="cometapi-test-key",
api_base=api_base,
)
request_body: Final = json.loads(route.calls.last.request.content)
assert str(route.calls.last.request.url) == f"{api_base}/chat/completions"
assert request_body["messages"] == [{"role": "user", "content": "Say OK"}]
assert response.choices[0].message.content == "OK"
# Integration test example (requires real API key)

75
tests/unit/llms/databricks/test_databricks.py
View file

@ -1,11 +1,13 @@
import asyncio
import json
import sys
from typing import Any, Dict, List
from types import ModuleType, SimpleNamespace
from typing import Any, Dict, Final, List
from unittest.mock import MagicMock, Mock, patch
import httpx
import pytest
import respx
import litellm
from litellm._version import version
@ -21,6 +23,21 @@ except ImportError:
databricks_sdk_installed = False
def _install_databricks_sdk(monkeypatch: pytest.MonkeyPatch, host: str, token: str) -> None:
config: Final = SimpleNamespace(
host=host,
authenticate=lambda: {"Authorization": f"Bearer {token}"},
)
sdk_module: Final = ModuleType("databricks.sdk")
package_module: Final = ModuleType("databricks")
package_module.__path__ = []
setattr(sdk_module, "WorkspaceClient", lambda: SimpleNamespace(config=config))
setattr(sdk_module, "useragent", SimpleNamespace(with_partner=lambda _: None))
setattr(package_module, "sdk", sdk_module)
monkeypatch.setitem(sys.modules, "databricks", package_module)
monkeypatch.setitem(sys.modules, "databricks.sdk", sdk_module)
def mock_chat_response() -> Dict[str, Any]:
return {
"id": "chatcmpl_3f78f09a-489c-4b8d-a587-f162c7497891",
@ -536,6 +553,62 @@ def test_embeddings_uses_databricks_sdk_if_api_key_and_base_not_specified(monkey
)
@pytest.mark.respx(assert_all_called=True)
def test_chat_completion_uses_databricks_sdk_credentials(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.delenv("DATABRICKS_API_BASE", raising=False)
monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
host: Final = "https://my.workspace.cloud.databricks.com"
token: Final = "sdk-test-token"
_install_databricks_sdk(monkeypatch, host, token)
messages: Final = [{"role": "user", "content": "How are you?"}]
route: Final = respx_mock.post(f"{host}/serving-endpoints/chat/completions").mock(
return_value=httpx.Response(200, json=mock_chat_response())
)
response: Final = litellm.completion(
model="databricks/databricks-gpt-5-6-sol",
messages=messages,
temperature=0.5,
)
request_body: Final = json.loads(route.calls.last.request.content)
assert str(route.calls.last.request.url) == f"{host}/serving-endpoints/chat/completions"
assert route.calls.last.request.headers["Authorization"] == f"Bearer {token}"
assert request_body["messages"] == messages
assert request_body["temperature"] == 0.5
assert response.choices[0].message.content == mock_chat_response()["choices"][0]["message"]["content"]
@pytest.mark.respx(assert_all_called=True)
def test_embedding_uses_databricks_sdk_credentials(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
) -> None:
monkeypatch.delenv("DATABRICKS_API_BASE", raising=False)
monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
host: Final = "https://my.workspace.cloud.databricks.com"
token: Final = "sdk-test-token"
_install_databricks_sdk(monkeypatch, host, token)
inputs: Final = ["hello"]
route: Final = respx_mock.post(f"{host}/serving-endpoints/embeddings").mock(
return_value=httpx.Response(200, json=mock_embedding_response())
)
response: Final = litellm.embedding(
model="databricks/databricks-qwen3-embedding-0-6b",
input=inputs,
)
request_body: Final = json.loads(route.calls.last.request.content)
assert str(route.calls.last.request.url) == f"{host}/serving-endpoints/embeddings"
assert route.calls.last.request.headers["Authorization"] == f"Bearer {token}"
assert request_body["input"] == inputs
assert response.to_dict()["data"] == mock_embedding_response()["data"]
def test_completion_with_prompt_caching_anthropic_model(monkeypatch):

98
tests/unit/llms/xai/test_xai_chat_transformation.py
View file

@ -1,9 +1,11 @@
import json
import os
from typing import Final
from unittest.mock import Mock, patch
import httpx
import pytest
import respx
import litellm
from litellm.llms.xai.chat.transformation import (
@ -167,6 +169,102 @@ class TestXAIUsageNormalization:
assert usage["total_tokens"] == 200
@pytest.mark.respx(assert_all_called=True)
def test_completion_strips_message_names_from_request(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post(f"{XAI_API_BASE}/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-xai",
"object": "chat.completion",
"created": 1234567890,
"model": "grok-4.20-beta-latest",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "OK"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
},
)
)
response: Final = litellm.completion(
model="xai/grok-4.20-beta-latest",
messages=[
{"role": "system", "content": "Be concise", "name": "system_prompt"},
{"role": "user", "content": "Say OK", "name": "caller"},
],
api_key="xai-test-key",
api_base=XAI_API_BASE,
)
request_body: Final = json.loads(route.calls.last.request.content)
assert request_body["messages"] == [
{"role": "system", "content": "Be concise"},
{"role": "user", "content": "Say OK", "name": "caller"},
]
assert response.choices[0].message.content == "OK"
@pytest.mark.respx(assert_all_called=True)
def test_streaming_include_usage_surfaces_xai_final_usage_chunk(respx_mock: respx.MockRouter) -> None:
frames: Final = (
{
"id": "chatcmpl-xai-stream",
"object": "chat.completion.chunk",
"created": 1234567890,
"model": "grok-4.20-beta-latest",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}, "finish_reason": None}],
},
{
"id": "chatcmpl-xai-stream",
"object": "chat.completion.chunk",
"created": 1234567890,
"model": "grok-4.20-beta-latest",
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
{
"id": "chatcmpl-xai-stream",
"object": "chat.completion.chunk",
"created": 1234567890,
"model": "grok-4.20-beta-latest",
"choices": [],
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
},
)
sse_body: Final = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + "data: [DONE]\n\n"
respx_mock.post(f"{XAI_API_BASE}/chat/completions").mock(
return_value=httpx.Response(
200,
content=sse_body.encode(),
headers={"content-type": "text/event-stream"},
)
)
chunks: Final = list(
litellm.completion(
model="xai/grok-4.20-beta-latest",
messages=[{"role": "user", "content": "Say OK"}],
stream=True,
stream_options={"include_usage": True},
api_key="xai-test-key",
api_base=XAI_API_BASE,
)
)
usage_chunks: Final = [chunk for chunk in chunks if getattr(chunk, "usage", None) is not None]
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "OK"
assert len(usage_chunks) == 1
assert (
usage_chunks[0].usage.prompt_tokens,
usage_chunks[0].usage.completion_tokens,
usage_chunks[0].usage.total_tokens,
) == (2, 1, 3)
class TestXAIChatWebSearchBilling:
_TOOL_DETAILS = {
"web_search_calls": 3,