mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
test(model_management): drive /model/block and /model/unblock through response serialization
The route-level regression test returns a real prisma row from a mocked update and asserts both routes serialize it to a 200 with the toggled blocked flag, which is exactly the path that raised AttributeError before the validator guard. Also binds the loop variable in the e2e poll lambda (ruff B023).
This commit is contained in:
parent
26e71ddc54
commit
fd72ae830c
2 changed files with 70 additions and 2 deletions
|
|
@ -271,8 +271,8 @@ class TestModelRoutes:
|
|||
assert response.blocked is expected
|
||||
_ = _poll(
|
||||
client.proxy,
|
||||
lambda: True
|
||||
if _model_blocked_flag(client, model_id) is expected
|
||||
lambda want=expected: True
|
||||
if _model_blocked_flag(client, model_id) is want
|
||||
else None,
|
||||
f"/model/info never reported blocked={expected} for {model_name!r} "
|
||||
f"after /model/{action}",
|
||||
|
|
|
|||
|
|
@ -4466,3 +4466,71 @@ class TestEnforceRpmTpmOnModelAdd:
|
|||
_raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True)
|
||||
assert expected_missing in str(exc_info.value.message)
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
|
||||
class TestBlockModelResponseSerialization:
|
||||
"""POST /model/block and /model/unblock return the raw prisma row through this
|
||||
route's `LiteLLM_ProxyModelTable | None` response validation. The row is not a
|
||||
dict, so the dict-assuming before-validator used to raise AttributeError inside
|
||||
FastAPI's serialization layer: a 500 for the caller after the DB write already
|
||||
landed. The routes must serialize the row to a 200 with the updated blocked flag."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "blocked"), [("/model/block", True), ("/model/unblock", False)]
|
||||
)
|
||||
def test_block_routes_serialize_prisma_row_to_200(self, route, blocked):
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from prisma import models as prisma_models
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
written_at = datetime(2026, 8, 29, tzinfo=timezone.utc)
|
||||
row_fields = {
|
||||
"model_id": "m-block-1",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"litellm_params": json.dumps({"model": "openai/gpt-4o-mini", "api_key": "encrypted-value"}),
|
||||
"model_info": json.dumps({"id": "m-block-1"}),
|
||||
"created_at": written_at,
|
||||
"created_by": "admin",
|
||||
"updated_at": written_at,
|
||||
"updated_by": "admin",
|
||||
}
|
||||
existing_row = prisma_models.LiteLLM_ProxyModelTable(blocked=not blocked, **row_fields)
|
||||
updated_row = prisma_models.LiteLLM_ProxyModelTable(blocked=blocked, **row_fields)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
|
||||
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
|
||||
|
||||
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: admin
|
||||
try:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
"litellm.proxy.proxy_server.llm_router",
|
||||
MagicMock(**{"get_model_ids.return_value": ["m-block-1"]}),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
patch( # test-quality-ok: stubs the cache write so the test observes only response serialization
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
patch( # test-quality-ok: audit logging is a background side effect outside this test's contract
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
):
|
||||
client = TestClient(app)
|
||||
response = client.post(route, json={"model_id": "m-block-1"})
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
body = response.json()
|
||||
assert body["model_id"] == "m-block-1"
|
||||
assert body["blocked"] is blocked
|
||||
assert body["litellm_params"] == {"model": "openai/gpt-4o-mini", "api_key": "encrypted-value"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue