test(proxy): cover key-created owners and global alias redirects for disallowed_models

This commit is contained in:
Jesus Velarde 2026-11-05 19:07:05 -08:00
parent a9ee60c5e1
commit 6bf64109eb
2 changed files with 19 additions and 1 deletions

View file

@ -145,7 +145,13 @@ def test_a_key_owner_cannot_remove_disallowed_models_but_can_add_them(gateway: G
allowed: Final = scenario.model()
disallowed: Final = scenario.model()
owner: Final = scenario.user(user_role="internal_user")
key: Final = scenario.key(user_id=owner, disallowed_models=[disallowed])
owner_key: Final = scenario.key(user_id=owner)
created: Final = gateway.request(
"POST", "/key/generate", {"user_id": owner, "disallowed_models": [disallowed]}, key=owner_key
)
assert created.status_code == 200, f"{created.status_code} {created.text}"
key: Final = string_value(JSON_OBJECT.validate_json(created.content)["key"])
scenario.cleanups.callback(scenario.delete_key, key)
cleared: Final = gateway.request("POST", "/key/update", {"key": key, "disallowed_models": []}, key=key)
assert cleared.status_code == 403, f"{cleared.status_code} {cleared.text}"
assert "disallowed_models" in cleared.text, cleared.text

View file

@ -2,6 +2,8 @@ from typing import Final
import pytest
import litellm
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import can_key_call_model
from litellm.proxy.auth.key_disallowed_models import (
@ -40,6 +42,16 @@ def test_a_disallowed_model_is_refused_however_it_is_named(fields: dict[str, obj
assert "is not available for this API key" in _denied_model(error)
def test_a_key_alias_that_a_global_alias_redirects_to_a_disallowed_model_is_refused(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "model_alias_map", {"public-name": "gpt-denied"})
token: Final = _token(disallowed_models=("gpt-denied",), aliases={"short": "public-name"})
with pytest.raises(ProxyException) as error:
ensure_key_may_call_model("short", token, None)
_denied_model(error)
@pytest.mark.parametrize("disallowed", [None, (), ("gpt-other",)])
def test_a_model_outside_disallowed_models_is_not_refused(disallowed: tuple[str, ...] | None) -> None:
assert first_disallowed_model(["gpt-allowed"], _token(disallowed_models=disallowed), None) is None