From d857b1c2ea58bf7ca5cf3db6f97f7b39d79f86b4 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 22 Jun 2026 18:37:17 -0700 Subject: [PATCH] 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. --- .../customer_endpoints.py | 10 ++ .../customer_endpoints.py | 18 ++++ .../test_customer_endpoints.py | 97 +++++++++++++++++++ 3 files changed, 125 insertions(+) create mode 100644 litellm/types/proxy/management_endpoints/customer_endpoints.py diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 50a1bc23a6d..4b309f15659 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -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", diff --git a/litellm/types/proxy/management_endpoints/customer_endpoints.py b/litellm/types/proxy/management_endpoints/customer_endpoints.py new file mode 100644 index 00000000000..3aba8d2d211 --- /dev/null +++ b/litellm/types/proxy/management_endpoints/customer_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 6c5ccd3562f..5243d868a7c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -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": []}. 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": , "message": }. + 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