Spaces:
Configuration error
Configuration error
File size: 10,264 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 |
import pytest
import asyncio
import aiohttp
import json
from httpx import AsyncClient
from typing import Any, Optional, List, Literal
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 = {}
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 = {}
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
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-4"], "gpt-4", 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-4", True), # No model restrictions
([], "gpt-4", 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"] == "401"
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-4"]}, 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-4
await mock_chat_completion(session=session, key=key, model="openai/gpt-4")
# Should fail with gpt-3.5-turbo
with pytest.raises(Exception) as exc_info:
await mock_chat_completion(
session=session, key=key, model="openai/gpt-3.5-turbo"
)
_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-4")
await mock_chat_completion(
session=session, key=key, model="openai/gpt-3.5-turbo"
)
# Non-OpenAI model should still fail
with pytest.raises(Exception) 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-4"], "name": "test-team"},
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}, headers=headers
)
assert response.status_code == 200
key = response.json()["key"]
# Test initial access
async with aiohttp.ClientSession() as session:
# Should work with gpt-4
await mock_chat_completion(session=session, key=key, model="openai/gpt-4")
# Should fail with gpt-3.5-turbo
with pytest.raises(Exception) as exc_info:
await mock_chat_completion(
session=session, key=key, model="openai/gpt-3.5-turbo"
)
_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-4")
await mock_chat_completion(
session=session, key=key, model="openai/gpt-3.5-turbo"
)
# Non-OpenAI model should still fail
with pytest.raises(Exception) 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"] == "401"
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"]
|