add tests for edge cases

This commit is contained in:
Rohan Shetty 2026-04-18 20:02:09 -04:00
parent 2184a3dab1
commit 91e694c6c7

View file

@ -51,3 +51,37 @@ class TestGetModelsUrl:
def test_localhost_without_v1(self):
"""Localhost with port but no /v1 should get /v1/models appended."""
assert self._get_url_for_api_base("http://localhost:11434") == "http://localhost:11434/v1/models"
def test_default_api_base(self):
"""If api_base is None, it should default to OpenAI and append /v1/models."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"data": []}
with patch("litellm.module_level_client.get", return_value=mock_response) as mock_get:
OpenAIGPTConfig().get_models(api_key="fake-key", api_base=None)
assert mock_get.call_args.kwargs["url"] == "https://api.openai.com/v1/models"
@patch("litellm.llms.openai.chat.gpt_transformation.get_secret_str", return_value="default-key")
def test_default_api_key(self, mock_get_secret):
"""If api_key is None, it should fetch from secrets."""
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"data": []}
with patch("litellm.module_level_client.get", return_value=mock_response) as mock_get:
OpenAIGPTConfig().get_models(api_key=None, api_base="https://example.com")
assert mock_get.call_args.kwargs["url"] == "https://example.com/v1/models"
assert mock_get.call_args.kwargs["headers"]["Authorization"] == "Bearer default-key"
def test_get_models_error(self):
"""If the API returns an error, get_models should raise an exception."""
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = "Error message"
import pytest
with patch("litellm.module_level_client.get", return_value=mock_response):
with pytest.raises(Exception) as exc:
OpenAIGPTConfig().get_models(api_key="fake-key", api_base="https://example.com")
assert "Failed to get models: Error message" in str(exc.value)