feat(proxy): type Customer Management response_model for OpenAPI coverage

Add response_model to the five remaining untyped /customer operations
(block, unblock, new, update, delete) so the generated OpenAPI schema
documents a concrete response body. new/update reuse the canonical
LiteLLM_EndUserTable (matching info/list); block, unblock, and delete
get small dedicated models in
litellm/types/proxy/management_endpoints/customer_endpoints.py.

Together with the already-typed info/list/daily-activity routes this
brings the Customer Management group to full response_model coverage.
Regression tests assert each public /customer/* route declares the
expected response_model and that /customer/new surfaces a typed schema
in app.openapi(), so dropping a response_model fails CI.
This commit is contained in:
ryan-crabbe-berri 2026-06-22 18:37:17 -07:00
parent 1cdb6cd3ac
commit d857b1c2ea
3 changed files with 125 additions and 0 deletions

View file

@ -32,6 +32,11 @@ from litellm.repositories.table_repositories import EndUserRepository
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from litellm.types.proxy.management_endpoints.customer_endpoints import (
BlockUsersResponse,
DeleteCustomersResponse,
UnblockUsersResponse,
)
router = APIRouter()
@ -46,6 +51,7 @@ router = APIRouter()
"/customer/block",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=BlockUsersResponse,
)
async def block_user(data: BlockUsers):
"""
@ -100,6 +106,7 @@ async def block_user(data: BlockUsers):
"/customer/unblock",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=UnblockUsersResponse,
)
async def unblock_user(data: BlockUsers):
"""
@ -222,6 +229,7 @@ async def _handle_customer_object_permission_update(
"/customer/new",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=LiteLLM_EndUserTable,
)
async def new_end_user(
data: NewCustomerRequest,
@ -489,6 +497,7 @@ async def end_user_info(
"/customer/update",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=LiteLLM_EndUserTable,
)
@router.post(
"/end_user/update",
@ -697,6 +706,7 @@ async def update_end_user(
"/customer/delete",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=DeleteCustomersResponse,
)
@router.post(
"/end_user/delete",

View file

@ -0,0 +1,18 @@
from typing import List
from pydantic import BaseModel
from litellm.models.end_user import LiteLLM_EndUserTable
class BlockUsersResponse(BaseModel):
blocked_users: List[LiteLLM_EndUserTable]
class UnblockUsersResponse(BaseModel):
blocked_users: List[str]
class DeleteCustomersResponse(BaseModel):
deleted_customers: int
message: str

View file

@ -1,8 +1,10 @@
from typing import List
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from fastapi.testclient import TestClient
from litellm.proxy._types import (
@ -13,6 +15,14 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.management_endpoints.customer_endpoints import router
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
from litellm.types.proxy.management_endpoints.customer_endpoints import (
BlockUsersResponse,
DeleteCustomersResponse,
UnblockUsersResponse,
)
app = FastAPI()
@ -317,6 +327,93 @@ def test_customer_endpoints_error_schema_consistency(
assert isinstance(error2[key], str), f"error2[{key}] should be a string"
EXPECTED_RESPONSE_MODELS = {
"/customer/block": BlockUsersResponse,
"/customer/unblock": UnblockUsersResponse,
"/customer/new": LiteLLM_EndUserTable,
"/customer/update": LiteLLM_EndUserTable,
"/customer/delete": DeleteCustomersResponse,
"/customer/info": LiteLLM_EndUserTable,
"/customer/list": List[LiteLLM_EndUserTable],
"/customer/daily/activity": SpendAnalyticsPaginatedResponse,
}
@pytest.mark.parametrize("path, expected_model", EXPECTED_RESPONSE_MODELS.items())
def test_customer_routes_declare_response_model(path, expected_model):
"""
Every public /customer/* operation must declare a typed response_model so
the generated OpenAPI schema documents the response body. Regression for the
OpenAPI response-type coverage goal: drop a response_model and this fails.
"""
route = next(r for r in router.routes if isinstance(r, APIRoute) and r.path == path)
assert route.response_model == expected_model
def test_customer_new_documented_in_openapi_schema():
"""
The response_model must surface in the OpenAPI schema as a concrete ref, not
an empty/default response. This is what the coverage metric measures.
"""
schema = app.openapi()["paths"]["/customer/new"]["post"]
json_schema = schema["responses"]["200"]["content"]["application/json"]["schema"]
assert json_schema["$ref"].endswith("/LiteLLM_EndUserTable")
def test_block_customer_success_serializes_through_response_model(
mock_prisma_client, mock_user_api_key_auth
):
"""
/customer/block returns {"blocked_users": [<end user rows>]}. With
response_model=BlockUsersResponse, a shape mismatch would raise a 500
ResponseValidationError, so a clean 200 proves the model matches runtime output.
"""
blocked_row = LiteLLM_EndUserTable(user_id="blocked-1", blocked=True)
mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(
return_value=blocked_row
)
response = client.post(
"/customer/block",
json={"user_ids": ["blocked-1"]},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
body = response.json()
assert body["blocked_users"][0]["user_id"] == "blocked-1"
assert body["blocked_users"][0]["blocked"] is True
def test_delete_customer_success_serializes_through_response_model(
mock_prisma_client, mock_user_api_key_auth
):
"""
/customer/delete returns {"deleted_customers": <int>, "message": <str>}.
response_model=DeleteCustomersResponse enforces that exact shape.
"""
existing = [
LiteLLM_EndUserTable(user_id="u1", blocked=False),
LiteLLM_EndUserTable(user_id="u2", blocked=False),
]
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
return_value=existing
)
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=2)
response = client.post(
"/customer/delete",
json={"user_ids": ["u1", "u2"]},
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
assert response.json() == {
"deleted_customers": 2,
"message": "Successfully deleted customers with ids: ['u1', 'u2']",
}
@pytest.mark.asyncio
async def test_get_customer_daily_activity_admin_param_passing(monkeypatch):
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth