test(custom_openai): cover model discovery branches

This commit is contained in:
Aditya Aggarwal 2026-07-26 16:09:54 +05:30
parent cf1445448b
commit db61d672b1
2 changed files with 57 additions and 37 deletions

View file

@ -1656,42 +1656,6 @@ def test_get_valid_models_openai_proxy(monkeypatch):
assert "litellm_proxy/gpt-5.5" in valid_models
def test_get_valid_models_custom_openai(monkeypatch):
from litellm.utils import get_valid_models
import litellm
mock_response_data = {
"object": "list",
"data": [
{
"id": "my-custom-model",
"object": "model",
"created": 1686935002,
"owned_by": "organization-owner",
},
],
}
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = mock_response_data
with patch.object(litellm.module_level_client, "get", return_value=mock_response):
valid_models = get_valid_models(
check_provider_endpoint=True,
custom_llm_provider="custom_openai",
api_key="sk-1234",
api_base="https://my-openai-compatible-endpoint/v1",
)
assert "custom_openai/my-custom-model" in valid_models
def test_custom_openai_get_models_requires_api_base():
from litellm.llms.openai_like.chat.transformation import CustomOpenAIChatConfig
with pytest.raises(ValueError):
CustomOpenAIChatConfig().get_models(api_base=None)
def test_get_valid_models_fireworks_ai(monkeypatch):
from litellm.utils import get_valid_models

View file

@ -1,5 +1,14 @@
from unittest.mock import MagicMock, patch
import litellm
import pytest
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
from litellm.llms.openai_like.chat.transformation import (
CustomOpenAIChatConfig,
OpenAILikeChatConfig,
)
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager, get_valid_models
def test_sanitize_usage_obj_handles_null_tokens():
@ -47,3 +56,50 @@ def test_sanitize_usage_obj_valid_usage():
# Assert
assert sanitized_json == original_json # The object should be unchanged
def test_get_valid_models_custom_openai():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"object": "list",
"data": [
{
"id": "my-custom-model",
"object": "model",
"created": 1686935002,
"owned_by": "organization-owner",
},
],
}
with patch.object(litellm.module_level_client, "get", return_value=mock_response) as mock_get:
valid_models = get_valid_models(
check_provider_endpoint=True,
custom_llm_provider="custom_openai",
api_key="sk-1234",
api_base="https://my-openai-compatible-endpoint/v1",
)
assert valid_models == ["custom_openai/my-custom-model"]
mock_get.assert_called_once_with(
url="https://my-openai-compatible-endpoint/v1/models",
headers={"Authorization": "Bearer sk-1234"},
)
def test_custom_openai_get_models_requires_api_base():
with pytest.raises(
ValueError,
match="api_base must be set to discover models for the custom_openai provider",
):
CustomOpenAIChatConfig().get_models(api_base=None)
def test_openai_model_info_uses_openai_config():
config = ProviderConfigManager.get_provider_model_info(
model=None,
provider=LlmProviders.OPENAI,
)
assert type(config) is OpenAIGPTConfig