88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for endpoint restrictions logic."""
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from vllm_omni.config.endpoint_policy import (
|
|
EndpointRestriction,
|
|
OmniServingCapability,
|
|
shutdown_unsupported_routes,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
SERVER_ROUTE = "/v1/completions"
|
|
REJECTION_REASON = "This route is banned!"
|
|
|
|
|
|
def _make_app_with_server_route():
|
|
"""Create a FastAPI app with a route to reject."""
|
|
app = FastAPI()
|
|
|
|
@app.post(SERVER_ROUTE)
|
|
async def existing_handler():
|
|
"""Returns a response; we want to disable this."""
|
|
return {"ok": True}
|
|
|
|
return app
|
|
|
|
|
|
def test_restricted_completions_returns_400():
|
|
"""Ensure that simple route rejection works."""
|
|
app = _make_app_with_server_route()
|
|
restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
|
|
shutdown_unsupported_routes(app, restrictions)
|
|
|
|
client = TestClient(app)
|
|
resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
|
|
assert resp.status_code == 400
|
|
body = resp.json()
|
|
assert body["error"]["message"] == REJECTION_REASON
|
|
assert body["error"]["type"] == "BadRequestError"
|
|
|
|
|
|
def test_unrestricted_completions_not_blocked():
|
|
"""Ensure that if we don't shutdown any routes, the route stays."""
|
|
app = _make_app_with_server_route()
|
|
shutdown_unsupported_routes(app, ())
|
|
|
|
client = TestClient(app)
|
|
resp = client.post(SERVER_ROUTE)
|
|
assert resp.status_code == 200
|
|
assert resp.json() == {"ok": True}
|
|
|
|
|
|
def test_route_rejection_is_idempotent():
|
|
"""Ensure that calling rejection twice doesn't break things."""
|
|
app = _make_app_with_server_route()
|
|
restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
|
|
shutdown_unsupported_routes(app, restrictions)
|
|
shutdown_unsupported_routes(app, restrictions)
|
|
|
|
client = TestClient(app)
|
|
resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
|
|
assert resp.status_code == 400
|
|
body = resp.json()
|
|
assert body["error"]["message"] == REJECTION_REASON
|
|
assert body["error"]["type"] == "BadRequestError"
|
|
|
|
|
|
def test_invalid_route_rejection():
|
|
"""Ensure that a restriction won't crash if a route doesn't exist."""
|
|
# NOTE: This test is for safety since the route organization is still messy.
|
|
# Since the restriction needs a valid OmniServingCapability, its implied the route
|
|
# is a valid vLLM route, but was already removed from the app by earlier processing.
|
|
app = FastAPI()
|
|
restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
|
|
shutdown_unsupported_routes(app, restrictions)
|
|
|
|
client = TestClient(app)
|
|
resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
|
|
assert resp.status_code == 400
|
|
body = resp.json()
|
|
assert body["error"]["message"] == REJECTION_REASON
|
|
assert body["error"]["type"] == "BadRequestError"
|