mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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:
parent
1cdb6cd3ac
commit
d857b1c2ea
3 changed files with 125 additions and 0 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue