scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Fix_mapped tests part 2

Browse source
This commit is contained in:
Sameer Kankute 2026-02-26 12:42:27 +05:30
parent 9c7f8138e1
commit 0debe92605
6 changed files with 69 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -1479,10 +1479,11 @@ class TestForwardHeaders:
# Create a mock request with custom headers
mock_request = MagicMock(spec=Request)
mock_request.state = None # Prevent MagicMock from returning a truthy _cached_headers
mock_request.method = "POST"
mock_request.url = MagicMock()
mock_request.url.path = "/test/endpoint"
# User headers that should be forwarded
user_headers = {
"x-custom-header": "custom-value",

98
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -1650,58 +1650,64 @@ async def test_global_spend_keys_endpoint_limit_validation(client, monkeypatch):
# Create a simple mock for prisma client with empty response
mock_prisma_client = MagicMock()
mock_db = MagicMock()
mock_query_raw = MagicMock()
mock_query_raw.return_value = asyncio.Future()
mock_query_raw.return_value.set_result([])
mock_query_raw = AsyncMock(return_value=[])
mock_db.query_raw = mock_query_raw
mock_prisma_client.db = mock_db
# Apply the mock to the prisma_client module
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
# Call the endpoint without specifying a limit
no_limit_response = client.get("/global/spend/keys")
assert no_limit_response.status_code == 200
mock_query_raw.assert_called_once_with('SELECT * FROM "Last30dKeysBySpend";')
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with valid input
normal_limit = "10"
good_input_response = client.get(f"/global/spend/keys?limit={normal_limit}")
assert good_input_response.status_code == 200
# Verify the mock was called with the correct parameters
mock_query_raw.assert_called_once_with(
'SELECT * FROM "Last30dKeysBySpend" LIMIT $1 ;', 10
# Override auth to bypass API key validation
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
)
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with SQL injection payload
sql_injection_limit = "10; DROP TABLE spend_logs; --"
response = client.get(f"/global/spend/keys?limit={sql_injection_limit}")
# Verify the response is a validation error (422)
assert response.status_code == 422
# Verify the mock was not called with the SQL injection payload
# This confirms that the validation happens before the database query
mock_query_raw.assert_not_called()
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with non-numeric input
non_numeric_limit = "abc"
response = client.get(f"/global/spend/keys?limit={non_numeric_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
# Test with negative number
negative_limit = "-5"
response = client.get(f"/global/spend/keys?limit={negative_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
# Test with zero
zero_limit = "0"
response = client.get(f"/global/spend/keys?limit={zero_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
try:
# Call the endpoint without specifying a limit
no_limit_response = client.get("/global/spend/keys")
assert no_limit_response.status_code == 200
mock_query_raw.assert_called_once_with('SELECT * FROM "Last30dKeysBySpend";')
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with valid input
normal_limit = "10"
good_input_response = client.get(f"/global/spend/keys?limit={normal_limit}")
assert good_input_response.status_code == 200
# Verify the mock was called with the correct parameters
mock_query_raw.assert_called_once_with(
'SELECT * FROM "Last30dKeysBySpend" LIMIT $1 ;', 10
)
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with SQL injection payload
sql_injection_limit = "10; DROP TABLE spend_logs; --"
response = client.get(f"/global/spend/keys?limit={sql_injection_limit}")
# Verify the response is a validation error (422)
assert response.status_code == 422
# Verify the mock was not called with the SQL injection payload
# This confirms that the validation happens before the database query
mock_query_raw.assert_not_called()
# Reset the mock for the next test
mock_query_raw.reset_mock()
# Test with non-numeric input
non_numeric_limit = "abc"
response = client.get(f"/global/spend/keys?limit={non_numeric_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
# Test with negative number
negative_limit = "-5"
response = client.get(f"/global/spend/keys?limit={negative_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
# Test with zero
zero_limit = "0"
response = client.get(f"/global/spend/keys?limit={zero_limit}")
assert response.status_code == 422
mock_query_raw.assert_not_called()
mock_query_raw.reset_mock()
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
@pytest.mark.asyncio

2
tests/test_litellm/proxy/test_health_check_functions.py
View file

@ -480,7 +480,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c
healthy = [{"model": "gpt-4"}]
unhealthy = []
async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None):
async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None):
return healthy, unhealthy
with patch(

9
tests/test_litellm/proxy/test_model_dump_with_preserved_fields.py
View file

@ -242,7 +242,7 @@ def test_full_output_structure_non_streaming():
)
result = model_dump_with_preserved_fields(response, exclude_unset=True)
# Top-level keys
# Top-level keys (usage is None when not explicitly set and excluded by exclude_unset=True)
assert set(result.keys()) == {
"id",
"choices",
@ -250,7 +250,6 @@ def test_full_output_structure_non_streaming():
"model",
"object",
"system_fingerprint",
"usage",
}
assert result["object"] == "chat.completion"
assert result["model"] == "gpt-4.1"
@ -270,12 +269,6 @@ def test_full_output_structure_non_streaming():
assert msg["content"] == "Hello!"
assert msg["role"] == "assistant"
# Usage structure
usage = result["usage"]
assert "prompt_tokens" in usage
assert "completion_tokens" in usage
assert "total_tokens" in usage
def test_full_output_structure_tool_calls():
"""

8
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -331,7 +331,8 @@ class TestProxyInitializationHelpers:
@patch("uvicorn.run")
@patch("builtins.print")
def test_max_requests_before_restart_flag(self, mock_print, mock_uvicorn_run):
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_max_requests_before_restart_flag(self, mock_setup_db, mock_print, mock_uvicorn_run):
"""Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests"""
from click.testing import CliRunner
@ -344,7 +345,10 @@ class TestProxyInitializationHelpers:
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with patch.dict(
os.environ, clean_env, clear=True,
), patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
@ -367,7 +371,7 @@ class TestProxyInitializationHelpers:
run_server, ["--local", "--max_requests_before_restart", "123"]
)
assert result.exit_code == 0
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_uvicorn_run.assert_called_once()
# Check that uvicorn.run was called with limit_max_requests parameter

11
tests/test_litellm/proxy/test_shared_health_check.py
View file

@ -1,10 +1,13 @@
import asyncio
import json
import pytest
import time
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
import pytest
from litellm.proxy.health_check_utils.shared_health_check_manager import (
SharedHealthCheckManager,
)
class TestSharedHealthCheckManager:
@ -272,7 +275,7 @@ class TestSharedHealthCheckManager:
)
# Should call perform_health_check and cache results
mock_perform.assert_called_once_with(model_list=model_list, details=True)
mock_perform.assert_called_once_with(model_list=model_list, details=True, max_concurrency=None)
assert healthy == expected_healthy
assert unhealthy == expected_unhealthy
@ -329,7 +332,7 @@ class TestSharedHealthCheckManager:
# Should fall back to local health check
mock_sleep.assert_called_once_with(2)
mock_perform.assert_called_once_with(model_list=model_list, details=True)
mock_perform.assert_called_once_with(model_list=model_list, details=True, max_concurrency=None)
assert healthy == expected_healthy
assert unhealthy == expected_unhealthy