mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
b76def0e5d
`pytest.raises(Exception)` with no `match=` passes on any error that broad. A TypeError from a refactor, a botched fixture, an import that moved: all of them read as the rejection the test claims to police, so the test goes green for the wrong reason and stays green after the behaviour it guards is gone. PT011 closes that gap for the 317 sites B017 could not reach, because B017 only fires on a single-statement body with no `as e` binding. Each pattern here is the message the code actually raised, recorded by running the sites under a plugin that logged the concrete type and text per call site, so the assertions describe observed behaviour rather than a guess. Where a site raises more than one message across its parametrize cases, the pattern is an alternation of what was seen; where the exception carries an empty `str()` and puts the text on `.message`, the site keeps a narrow `noqa` with the reason. PT014 removes four parametrize cases that were listed twice. The duplicate re-runs an assertion that already passed, and it usually marks a case someone meant to vary and forgot to edit.
305 lines
11 KiB
Python
305 lines
11 KiB
Python
import pytest
|
|
import asyncio
|
|
import aiohttp
|
|
import json
|
|
from httpx import AsyncClient
|
|
from openai import PermissionDeniedError
|
|
from typing import Any, Optional, List, Literal
|
|
|
|
|
|
# The proxy strips client-supplied `mock_response` unless the calling key or
|
|
# team has this admin-metadata flag set. See `_UNTRUSTED_ROOT_CONTROL_FIELDS`
|
|
# in litellm/proxy/litellm_pre_call_utils.py.
|
|
_ALLOW_CLIENT_MOCK_METADATA = {"allow_client_mock_response": True}
|
|
|
|
|
|
async def generate_key(
|
|
session, models: Optional[List[str]] = None, team_id: Optional[str] = None
|
|
):
|
|
"""Helper function to generate a key with specific model access controls"""
|
|
url = "http://0.0.0.0:4000/key/generate"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
|
|
if models is not None:
|
|
data["models"] = models
|
|
if team_id is not None:
|
|
data["team_id"] = team_id
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
return await response.json()
|
|
|
|
|
|
async def generate_team(session, models: Optional[List[str]] = None):
|
|
"""Helper function to generate a team with specific model access"""
|
|
url = "http://0.0.0.0:4000/team/new"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)}
|
|
if models is not None:
|
|
data["models"] = models
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
return await response.json()
|
|
|
|
|
|
async def mock_chat_completion(session, key: str, model: str):
|
|
"""Make a chat completion request using OpenAI SDK"""
|
|
from openai import AsyncOpenAI
|
|
from litellm._uuid import uuid
|
|
|
|
client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000/v1")
|
|
|
|
response = await client.chat.completions.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": f"Say hello! {uuid.uuid4()}"}],
|
|
extra_body={
|
|
"mock_response": "mock_response",
|
|
},
|
|
)
|
|
return response
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"key_models, test_model, expect_success",
|
|
[
|
|
(["openai/*"], "anthropic/claude-2", False), # Non-matching model
|
|
(["gpt-5.5"], "gpt-5.5", True), # Exact model match
|
|
(["bedrock/*"], "bedrock/anthropic.claude-3", True), # Bedrock wildcard
|
|
(["bedrock/anthropic.*"], "bedrock/anthropic.claude-3", True), # Pattern match
|
|
(["bedrock/anthropic.*"], "bedrock/amazon.titan", False), # Pattern non-match
|
|
(None, "gpt-5.5", True), # No model restrictions
|
|
([], "gpt-5.5", True), # Empty model list
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_model_access_patterns(key_models, test_model, expect_success):
|
|
"""
|
|
Test model access patterns for API keys:
|
|
1. Create key with specific model access pattern
|
|
2. Attempt to make completion with test model
|
|
3. Verify access is granted/denied as expected
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
# Generate key with specified model access
|
|
key_gen = await generate_key(session=session, models=key_models)
|
|
key = key_gen["key"]
|
|
|
|
try:
|
|
response = await mock_chat_completion(
|
|
session=session,
|
|
key=key,
|
|
model=test_model,
|
|
)
|
|
if not expect_success:
|
|
pytest.fail(f"Expected request to fail for model {test_model}")
|
|
assert (
|
|
response is not None
|
|
), "Should get valid response when access is allowed"
|
|
except Exception as e:
|
|
if expect_success:
|
|
pytest.fail(f"Expected request to succeed but got error: {e}")
|
|
_error_body = e.body
|
|
|
|
# Assert error structure and values
|
|
assert _error_body["type"] == "key_model_access_denied"
|
|
assert _error_body["param"] == "model"
|
|
assert _error_body["code"] == "403"
|
|
assert "key not allowed to access model" in _error_body["message"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_access_update():
|
|
"""
|
|
Test updating model access for an existing key:
|
|
1. Create key with restricted model access
|
|
2. Verify access patterns
|
|
3. Update key with new model access
|
|
4. Verify new access patterns
|
|
"""
|
|
client = AsyncClient(base_url="http://0.0.0.0:4000")
|
|
headers = {"Authorization": "Bearer sk-1234"}
|
|
|
|
# Create initial key with restricted access
|
|
response = await client.post(
|
|
"/key/generate",
|
|
json={
|
|
"models": ["openai/gpt-5.5"],
|
|
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200
|
|
key_data = response.json()
|
|
key = key_data["key"]
|
|
|
|
# Test initial access
|
|
async with aiohttp.ClientSession() as session:
|
|
# Should work with gpt-5.5
|
|
await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5")
|
|
|
|
# Should fail with gpt-5-mini
|
|
with pytest.raises(PermissionDeniedError) as exc_info:
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="openai/gpt-5-mini"
|
|
)
|
|
_validate_model_access_exception(
|
|
exc_info.value, expected_type="key_model_access_denied"
|
|
)
|
|
|
|
# Update key with new model access
|
|
response = await client.post(
|
|
"/key/update", json={"key": key, "models": ["openai/*"]}, headers=headers
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Test updated access
|
|
async with aiohttp.ClientSession() as session:
|
|
# Both models should now work
|
|
await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5")
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="openai/gpt-5-mini"
|
|
)
|
|
|
|
# Non-OpenAI model should still fail
|
|
with pytest.raises(PermissionDeniedError) as exc_info:
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="anthropic/claude-2"
|
|
)
|
|
_validate_model_access_exception(
|
|
exc_info.value, expected_type="key_model_access_denied"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"team_models, test_model, expect_success",
|
|
[
|
|
(["openai/*"], "anthropic/claude-2", False), # Non-matching model
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_team_model_access_patterns(team_models, test_model, expect_success):
|
|
"""
|
|
Test model access patterns for team-based API keys:
|
|
1. Create team with specific model access pattern
|
|
2. Generate key for that team
|
|
3. Attempt to make completion with test model
|
|
4. Verify access is granted/denied as expected
|
|
"""
|
|
client = AsyncClient(base_url="http://0.0.0.0:4000")
|
|
headers = {"Authorization": "Bearer sk-1234"}
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
try:
|
|
team_gen = await generate_team(session=session, models=team_models)
|
|
print("created team", team_gen)
|
|
team_id = team_gen["team_id"]
|
|
key_gen = await generate_key(session=session, team_id=team_id)
|
|
print("created key", key_gen)
|
|
key = key_gen["key"]
|
|
response = await mock_chat_completion(
|
|
session=session,
|
|
key=key,
|
|
model=test_model,
|
|
)
|
|
if not expect_success:
|
|
pytest.fail(f"Expected request to fail for model {test_model}")
|
|
assert (
|
|
response is not None
|
|
), "Should get valid response when access is allowed"
|
|
except Exception as e:
|
|
if expect_success:
|
|
pytest.fail(f"Expected request to succeed but got error: {e}")
|
|
_validate_model_access_exception(
|
|
e, expected_type="team_model_access_denied"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_team_model_access_update():
|
|
"""
|
|
Test updating model access for a team:
|
|
1. Create team with restricted model access
|
|
2. Verify access patterns
|
|
3. Update team with new model access
|
|
4. Verify new access patterns
|
|
"""
|
|
client = AsyncClient(base_url="http://0.0.0.0:4000")
|
|
headers = {"Authorization": "Bearer sk-1234"}
|
|
|
|
# Create initial team with restricted access
|
|
response = await client.post(
|
|
"/team/new",
|
|
json={
|
|
"models": ["openai/gpt-5.5"],
|
|
"name": "test-team",
|
|
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200
|
|
team_data = response.json()
|
|
team_id = team_data["team_id"]
|
|
|
|
# Generate a key for this team
|
|
response = await client.post(
|
|
"/key/generate",
|
|
json={
|
|
"team_id": team_id,
|
|
"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA),
|
|
},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200
|
|
key = response.json()["key"]
|
|
|
|
# Test initial access
|
|
async with aiohttp.ClientSession() as session:
|
|
# Should work with gpt-5.5
|
|
await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5")
|
|
|
|
# Should fail with gpt-5-mini
|
|
with pytest.raises(PermissionDeniedError) as exc_info:
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="openai/gpt-5-mini"
|
|
)
|
|
_validate_model_access_exception(
|
|
exc_info.value, expected_type="team_model_access_denied"
|
|
)
|
|
|
|
# Update team with new model access
|
|
response = await client.post(
|
|
"/team/update",
|
|
json={"team_id": team_id, "models": ["openai/*"]},
|
|
headers=headers,
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Test updated access
|
|
async with aiohttp.ClientSession() as session:
|
|
# Both models should now work
|
|
await mock_chat_completion(session=session, key=key, model="openai/gpt-5.5")
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="openai/gpt-5-mini"
|
|
)
|
|
|
|
# Non-OpenAI model should still fail
|
|
with pytest.raises(PermissionDeniedError) as exc_info:
|
|
await mock_chat_completion(
|
|
session=session, key=key, model="anthropic/claude-2"
|
|
)
|
|
_validate_model_access_exception(
|
|
exc_info.value, expected_type="team_model_access_denied"
|
|
)
|
|
|
|
|
|
def _validate_model_access_exception(
|
|
e: Exception,
|
|
expected_type: Literal["key_model_access_denied", "team_model_access_denied"],
|
|
):
|
|
_error_body = e.body
|
|
|
|
# Assert error structure and values
|
|
assert _error_body["type"] == expected_type
|
|
assert _error_body["param"] == "model"
|
|
assert _error_body["code"] == "403"
|
|
if expected_type == "key_model_access_denied":
|
|
assert "key not allowed to access model" in _error_body["message"]
|
|
elif expected_type == "team_model_access_denied":
|
|
assert "eam not allowed to access model" in _error_body["message"]
|