mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
test(proxy): cover key-created owners and global alias redirects for disallowed_models
This commit is contained in:
parent
a9ee60c5e1
commit
6bf64109eb
2 changed files with 19 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue