diff --git a/tests/integration/authorization/test_key_disallowed_models.py b/tests/integration/authorization/test_key_disallowed_models.py index 2a81876d453..2000fbfe8e1 100644 --- a/tests/integration/authorization/test_key_disallowed_models.py +++ b/tests/integration/authorization/test_key_disallowed_models.py @@ -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 diff --git a/tests/unit/proxy/auth/test_key_disallowed_models.py b/tests/unit/proxy/auth/test_key_disallowed_models.py index 0986b931464..b4bb90981e3 100644 --- a/tests/unit/proxy/auth/test_key_disallowed_models.py +++ b/tests/unit/proxy/auth/test_key_disallowed_models.py @@ -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