2024-02-16 08:55:21 -08:00
|
|
|
# What this tests ?
|
|
|
|
|
## Tests /team endpoints.
|
|
|
|
|
import pytest
|
|
|
|
|
import asyncio
|
|
|
|
|
import aiohttp
|
2024-02-21 16:53:12 -08:00
|
|
|
import time, uuid
|
2024-02-16 08:55:21 -08:00
|
|
|
from openai import AsyncOpenAI
|
2024-10-08 21:57:03 -07:00
|
|
|
from typing import Optional
|
2025-03-06 22:04:36 -08:00
|
|
|
import openai
|
|
|
|
|
from unittest.mock import MagicMock, patch
|
2024-10-08 21:57:03 -07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
async def get_user_info(session, get_user, call_user, view_all: Optional[bool] = None):
|
|
|
|
|
"""
|
|
|
|
|
Make sure only models user has access to are returned
|
|
|
|
|
"""
|
|
|
|
|
if view_all is True:
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/user/info"
|
2024-10-08 21:57:03 -07:00
|
|
|
else:
|
2025-12-18 13:16:15 -08:00
|
|
|
url = f"http://localhost:4000/user/info?user_id={get_user}"
|
2024-10-08 21:57:03 -07:00
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {call_user}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async with session.get(url, headers=headers) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
if call_user != get_user:
|
|
|
|
|
return status
|
|
|
|
|
else:
|
|
|
|
|
print(f"call_user: {call_user}; get_user: {get_user}")
|
|
|
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
|
|
|
return await response.json()
|
2024-02-16 08:55:21 -08:00
|
|
|
|
|
|
|
|
|
2025-12-18 13:16:15 -08:00
|
|
|
async def wait_for_team_member_spend_update(
|
|
|
|
|
session, user_id, team_id, expected_min_spend, max_wait=10
|
|
|
|
|
):
|
|
|
|
|
"""
|
|
|
|
|
Wait for the team member spend update to be committed to the database.
|
|
|
|
|
Polls the user info endpoint until the spend is updated.
|
|
|
|
|
This is needed because spend updates are queued asynchronously and committed periodically.
|
|
|
|
|
"""
|
|
|
|
|
start_time = time.time()
|
|
|
|
|
initial_spend = None
|
|
|
|
|
while time.time() - start_time < max_wait:
|
|
|
|
|
try:
|
|
|
|
|
user_info = await get_user_info(session, user_id, call_user="sk-1234")
|
|
|
|
|
if user_info.get("teams"):
|
|
|
|
|
for team in user_info["teams"]:
|
|
|
|
|
if team.get("team_id") == team_id:
|
|
|
|
|
for membership in team.get("team_memberships", []):
|
|
|
|
|
spend = membership.get("spend", 0.0)
|
|
|
|
|
if initial_spend is None:
|
|
|
|
|
initial_spend = spend
|
|
|
|
|
print(f"Initial team member spend: {spend}")
|
2026-03-13 12:35:20 -07:00
|
|
|
|
2025-12-18 13:16:15 -08:00
|
|
|
if spend >= expected_min_spend:
|
|
|
|
|
print(f"[OK] Team member spend updated: {spend} >= {expected_min_spend}")
|
|
|
|
|
return True
|
2026-03-13 12:35:20 -07:00
|
|
|
|
|
|
|
|
print(f"[WAITING] Team member spend: {spend}, expected >= {expected_min_spend}, elapsed: {time.time() - start_time:.1f}s")
|
2025-12-18 13:16:15 -08:00
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
print(f"Error checking team member spend: {e}")
|
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
|
print(f"[TIMEOUT] Timeout waiting for team member spend update (expected >= {expected_min_spend})")
|
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
2024-03-07 07:56:51 -08:00
|
|
|
async def new_user(
|
2024-05-22 19:19:51 -07:00
|
|
|
session,
|
|
|
|
|
i,
|
|
|
|
|
user_id=None,
|
|
|
|
|
budget=None,
|
|
|
|
|
budget_duration=None,
|
|
|
|
|
models=["azure-models"],
|
|
|
|
|
team_id=None,
|
2024-06-08 16:27:14 -07:00
|
|
|
user_email=None,
|
2024-03-07 07:56:51 -08:00
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/user/new"
|
2024-02-21 16:53:12 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
data = {
|
2024-03-07 07:56:51 -08:00
|
|
|
"models": models,
|
2024-02-21 16:53:12 -08:00
|
|
|
"aliases": {"mistral-7b": "gpt-3.5-turbo"},
|
|
|
|
|
"duration": None,
|
|
|
|
|
"max_budget": budget,
|
|
|
|
|
"budget_duration": budget_duration,
|
2024-06-08 16:27:14 -07:00
|
|
|
"user_email": user_email,
|
2024-02-21 16:53:12 -08:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if user_id is not None:
|
|
|
|
|
data["user_id"] = user_id
|
|
|
|
|
|
2024-05-22 19:19:51 -07:00
|
|
|
if team_id is not None:
|
|
|
|
|
data["team_id"] = team_id
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
2025-02-01 11:23:00 -08:00
|
|
|
raise Exception(
|
|
|
|
|
f"Request {i} did not return a 200 status code: {status}, response: {response_text}"
|
|
|
|
|
)
|
2024-02-21 16:53:12 -08:00
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-05-22 19:19:51 -07:00
|
|
|
async def add_member(
|
2024-06-08 19:03:45 -07:00
|
|
|
session, i, team_id, user_id=None, user_email=None, max_budget=None, members=None
|
2024-05-22 19:19:51 -07:00
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/member_add"
|
2024-04-15 10:29:21 -07:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
data = {"team_id": team_id, "member": {"role": "user"}}
|
|
|
|
|
if user_email is not None:
|
|
|
|
|
data["member"]["user_email"] = user_email
|
|
|
|
|
elif user_id is not None:
|
|
|
|
|
data["member"]["user_id"] = user_id
|
2024-06-08 19:03:45 -07:00
|
|
|
elif members is not None:
|
|
|
|
|
data["member"] = members
|
2024-04-15 10:29:21 -07:00
|
|
|
|
2024-05-22 19:19:51 -07:00
|
|
|
if max_budget is not None:
|
|
|
|
|
data["max_budget_in_team"] = max_budget
|
|
|
|
|
|
2024-06-08 19:03:45 -07:00
|
|
|
print("sent data: {}".format(data))
|
2024-04-15 10:29:21 -07:00
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"ADD MEMBER Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2025-02-01 11:23:00 -08:00
|
|
|
async def update_member(
|
|
|
|
|
session,
|
|
|
|
|
i,
|
|
|
|
|
team_id,
|
|
|
|
|
user_id=None,
|
|
|
|
|
user_email=None,
|
|
|
|
|
max_budget=None,
|
|
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/member_update"
|
2025-02-01 11:23:00 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
data = {"team_id": team_id}
|
|
|
|
|
if user_id is not None:
|
|
|
|
|
data["user_id"] = user_id
|
|
|
|
|
elif user_email is not None:
|
|
|
|
|
data["user_email"] = user_email
|
|
|
|
|
|
|
|
|
|
if max_budget is not None:
|
|
|
|
|
data["max_budget_in_team"] = max_budget
|
|
|
|
|
|
|
|
|
|
print("sent data: {}".format(data))
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"ADD MEMBER Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(
|
|
|
|
|
f"Request {i} did not return a 200 status code: {status}, response: {response_text}"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-06-08 16:27:14 -07:00
|
|
|
async def delete_member(session, i, team_id, user_id=None, user_email=None):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/member_delete"
|
2024-03-01 09:14:08 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
2024-06-08 16:27:14 -07:00
|
|
|
data = {"team_id": team_id}
|
|
|
|
|
if user_id is not None:
|
|
|
|
|
data["user_id"] = user_id
|
|
|
|
|
elif user_email is not None:
|
|
|
|
|
data["user_email"] = user_email
|
2024-03-01 09:14:08 -08:00
|
|
|
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
async def generate_key(
|
2024-02-16 08:55:21 -08:00
|
|
|
session,
|
|
|
|
|
i,
|
2024-02-21 16:53:12 -08:00
|
|
|
budget=None,
|
|
|
|
|
budget_duration=None,
|
|
|
|
|
models=["azure-models", "gpt-4", "dall-e-3"],
|
|
|
|
|
team_id=None,
|
2024-02-16 08:55:21 -08:00
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/key/generate"
|
2024-02-21 16:53:12 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
data = {
|
|
|
|
|
"models": models,
|
|
|
|
|
"duration": None,
|
|
|
|
|
"max_budget": budget,
|
|
|
|
|
"budget_duration": budget_duration,
|
|
|
|
|
}
|
|
|
|
|
if team_id is not None:
|
|
|
|
|
data["team_id"] = team_id
|
|
|
|
|
|
|
|
|
|
print(f"data: {data}")
|
|
|
|
|
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def chat_completion(session, key, model="gpt-4"):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/chat/completions"
|
2024-02-21 16:53:12 -08:00
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {key}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
|
|
|
|
data = {
|
|
|
|
|
"model": model,
|
|
|
|
|
"messages": [
|
|
|
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
|
|
|
{"role": "user", "content": "Hello!"},
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
for i in range(3):
|
|
|
|
|
try:
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(
|
|
|
|
|
f"Request did not return a 200 status code: {status}. Response: {response_text}"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
except Exception as e:
|
|
|
|
|
if "Request did not return a 200 status code" in str(e):
|
|
|
|
|
raise e
|
|
|
|
|
else:
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
|
2024-03-07 07:56:51 -08:00
|
|
|
async def new_team(session, i, user_id=None, member_list=None, model_aliases=None):
|
|
|
|
|
import json
|
|
|
|
|
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/new"
|
2024-02-16 08:55:21 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
2024-03-07 07:56:51 -08:00
|
|
|
data = {"team_alias": "my-new-team"}
|
2024-02-21 16:53:12 -08:00
|
|
|
if user_id is not None:
|
|
|
|
|
data["members_with_roles"] = [{"role": "user", "user_id": user_id}]
|
|
|
|
|
elif member_list is not None:
|
|
|
|
|
data["members_with_roles"] = member_list
|
|
|
|
|
|
2024-03-07 07:56:51 -08:00
|
|
|
if model_aliases is not None:
|
|
|
|
|
data["model_aliases"] = model_aliases
|
|
|
|
|
|
|
|
|
|
print(f"data: {data}")
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-03-28 15:59:35 -07:00
|
|
|
async def update_team(session, i, team_id, user_id=None, member_list=None, **kwargs):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/update"
|
2024-02-21 16:53:12 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
2024-03-28 15:59:35 -07:00
|
|
|
data = {"team_id": team_id, **kwargs}
|
2024-02-21 16:53:12 -08:00
|
|
|
if user_id is not None:
|
|
|
|
|
data["members_with_roles"] = [{"role": "user", "user_id": user_id}]
|
|
|
|
|
elif member_list is not None:
|
|
|
|
|
data["members_with_roles"] = member_list
|
|
|
|
|
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def delete_team(
|
|
|
|
|
session,
|
|
|
|
|
i,
|
|
|
|
|
team_id,
|
|
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/delete"
|
2024-02-21 16:53:12 -08:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
data = {
|
|
|
|
|
"team_ids": [team_id],
|
|
|
|
|
}
|
|
|
|
|
|
2024-02-16 08:55:21 -08:00
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
|
|
|
|
|
print(f"Response {i} (Status code: {status}):")
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-08-10 16:36:43 -07:00
|
|
|
async def list_teams(
|
|
|
|
|
session,
|
|
|
|
|
i,
|
|
|
|
|
):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = "http://localhost:4000/team/list"
|
2024-08-10 16:36:43 -07:00
|
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
|
|
|
|
|
|
async with session.get(url, headers=headers) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
|
|
|
|
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
2024-02-16 08:55:21 -08:00
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_team_new():
|
|
|
|
|
"""
|
|
|
|
|
Make 20 parallel calls to /user/new. Assert all worked.
|
|
|
|
|
"""
|
2024-02-21 16:53:12 -08:00
|
|
|
user_id = f"{uuid.uuid4()}"
|
2024-02-16 08:55:21 -08:00
|
|
|
async with aiohttp.ClientSession() as session:
|
2024-02-21 16:53:12 -08:00
|
|
|
new_user(session=session, i=0, user_id=user_id)
|
|
|
|
|
tasks = [new_team(session, i, user_id=user_id) for i in range(1, 11)]
|
2024-02-16 08:55:21 -08:00
|
|
|
await asyncio.gather(*tasks)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def get_team_info(session, get_team, call_key):
|
2025-12-18 13:16:15 -08:00
|
|
|
url = f"http://localhost:4000/team/info?team_id={get_team}"
|
2024-02-16 08:55:21 -08:00
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {call_key}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async with session.get(url, headers=headers) as response:
|
|
|
|
|
status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
print(response_text)
|
|
|
|
|
print()
|
|
|
|
|
|
2025-03-06 22:04:36 -08:00
|
|
|
if status == 404:
|
|
|
|
|
raise openai.NotFoundError(
|
|
|
|
|
message="404 received", response=MagicMock(), body=None
|
|
|
|
|
)
|
|
|
|
|
|
2024-02-16 08:55:21 -08:00
|
|
|
if status != 200:
|
|
|
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
|
|
|
return await response.json()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_team_info():
|
2024-04-16 11:39:52 -07:00
|
|
|
"""
|
|
|
|
|
Scenario 1:
|
|
|
|
|
- test with admin key -> expect to work
|
|
|
|
|
Scenario 2:
|
|
|
|
|
- test with team key -> expect to work
|
|
|
|
|
Scenario 3:
|
|
|
|
|
- test with non-team key -> expect to fail
|
|
|
|
|
"""
|
2024-02-16 08:55:21 -08:00
|
|
|
async with aiohttp.ClientSession() as session:
|
2024-04-16 11:39:52 -07:00
|
|
|
"""
|
|
|
|
|
Scenario 1 - as admin
|
|
|
|
|
"""
|
2024-02-16 08:55:21 -08:00
|
|
|
new_team_data = await new_team(
|
|
|
|
|
session,
|
|
|
|
|
0,
|
|
|
|
|
)
|
|
|
|
|
team_id = new_team_data["team_id"]
|
|
|
|
|
## as admin ##
|
|
|
|
|
await get_team_info(session=session, get_team=team_id, call_key="sk-1234")
|
2024-04-16 11:39:52 -07:00
|
|
|
"""
|
|
|
|
|
Scenario 2 - as team key
|
|
|
|
|
"""
|
|
|
|
|
key_gen = await generate_key(session=session, i=0, team_id=team_id)
|
|
|
|
|
key = key_gen["key"]
|
|
|
|
|
|
|
|
|
|
await get_team_info(session=session, get_team=team_id, call_key=key)
|
|
|
|
|
|
|
|
|
|
"""
|
|
|
|
|
Scenario 3 - as non-team key
|
|
|
|
|
"""
|
|
|
|
|
key_gen = await generate_key(session=session, i=0)
|
|
|
|
|
key = key_gen["key"]
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
await get_team_info(session=session, get_team=team_id, call_key=key)
|
2024-11-08 22:07:17 +05:30
|
|
|
pytest.fail("Expected call to fail")
|
2024-04-16 11:39:52 -07:00
|
|
|
except Exception as e:
|
|
|
|
|
pass
|
2024-02-21 16:53:12 -08:00
|
|
|
|
|
|
|
|
|
2024-04-15 10:29:21 -07:00
|
|
|
"""
|
|
|
|
|
- Create team
|
|
|
|
|
- Add user (user exists in db)
|
|
|
|
|
- Update team
|
|
|
|
|
- Check if it works
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
"""
|
|
|
|
|
- Create team
|
|
|
|
|
- Add user (user doesn't exist in db)
|
|
|
|
|
- Update team
|
|
|
|
|
- Check if it works
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
@pytest.mark.asyncio
|
2024-04-15 10:29:21 -07:00
|
|
|
async def test_team_update_sc_2():
|
2024-02-21 16:53:12 -08:00
|
|
|
"""
|
2024-04-15 10:29:21 -07:00
|
|
|
- Create team
|
2024-06-08 19:03:45 -07:00
|
|
|
- Add 3 users (doesn't exist in db)
|
2024-04-15 10:29:21 -07:00
|
|
|
- Change team alias
|
|
|
|
|
- Check if it works
|
|
|
|
|
- Assert team object unchanged besides team alias
|
2024-02-21 16:53:12 -08:00
|
|
|
"""
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
|
## Create admin
|
|
|
|
|
admin_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=admin_user)
|
|
|
|
|
## Create team with 1 admin and 1 user
|
|
|
|
|
member_list = [
|
|
|
|
|
{"role": "admin", "user_id": admin_user},
|
|
|
|
|
]
|
|
|
|
|
team_data = await new_team(session=session, i=0, member_list=member_list)
|
2024-06-08 19:03:45 -07:00
|
|
|
## Create 10 normal users
|
|
|
|
|
members = [
|
|
|
|
|
{"role": "user", "user_id": f"krrish_{uuid.uuid4()}@berri.ai"}
|
|
|
|
|
for _ in range(10)
|
|
|
|
|
]
|
2024-04-15 10:29:21 -07:00
|
|
|
await add_member(
|
2024-06-08 19:03:45 -07:00
|
|
|
session=session, i=0, team_id=team_data["team_id"], members=members
|
|
|
|
|
)
|
|
|
|
|
## ASSERT TEAM SIZE
|
|
|
|
|
team_info = await get_team_info(
|
|
|
|
|
session=session, get_team=team_data["team_id"], call_key="sk-1234"
|
2024-02-21 16:53:12 -08:00
|
|
|
)
|
|
|
|
|
|
2024-06-08 19:03:45 -07:00
|
|
|
assert len(team_info["team_info"]["members_with_roles"]) == 12
|
|
|
|
|
|
2024-04-15 10:29:21 -07:00
|
|
|
## CHANGE TEAM ALIAS
|
|
|
|
|
|
|
|
|
|
new_team_data = await update_team(
|
|
|
|
|
session=session, i=0, team_id=team_data["team_id"], team_alias="my-new-team"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert new_team_data["data"]["team_alias"] == "my-new-team"
|
|
|
|
|
print(f"team_data: {team_data}")
|
|
|
|
|
## assert rest of object is the same
|
|
|
|
|
for k, v in new_team_data["data"].items():
|
|
|
|
|
if (
|
|
|
|
|
k == "members_with_roles"
|
|
|
|
|
): # assert 1 more member (role: "user", user_email: $user_email)
|
|
|
|
|
len(new_team_data["data"][k]) == len(team_data[k]) + 1
|
|
|
|
|
elif (
|
|
|
|
|
k == "created_at"
|
|
|
|
|
or k == "updated_at"
|
|
|
|
|
or k == "model_spend"
|
|
|
|
|
or k == "model_max_budget"
|
|
|
|
|
or k == "model_id"
|
|
|
|
|
or k == "litellm_organization_table"
|
2025-05-09 09:38:03 -07:00
|
|
|
or k == "object_permission_id"
|
2025-05-09 11:59:13 -07:00
|
|
|
or k == "object_permission"
|
2024-04-15 10:29:21 -07:00
|
|
|
or k == "litellm_model_table"
|
2026-01-24 13:17:32 -08:00
|
|
|
or k == "policies"
|
2026-02-05 09:40:21 +05:30
|
|
|
or k == "allow_team_guardrail_config"
|
2026-02-21 10:24:53 -08:00
|
|
|
or k == "projects"
|
2024-04-15 10:29:21 -07:00
|
|
|
):
|
|
|
|
|
pass
|
|
|
|
|
else:
|
|
|
|
|
assert new_team_data["data"][k] == team_data[k]
|
2024-03-28 15:59:35 -07:00
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
|
2024-08-10 17:12:09 -07:00
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_team_member_add_email():
|
2025-05-27 17:44:13 -07:00
|
|
|
from tests.test_users import get_user_info
|
2024-08-10 17:12:09 -07:00
|
|
|
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
|
## Create admin
|
|
|
|
|
admin_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=admin_user)
|
|
|
|
|
## Create team with 1 admin and 1 user
|
|
|
|
|
member_list = [
|
|
|
|
|
{"role": "admin", "user_id": admin_user},
|
|
|
|
|
]
|
|
|
|
|
team_data = await new_team(session=session, i=0, member_list=member_list)
|
|
|
|
|
## Add 1 user via email
|
|
|
|
|
user_email = "krrish{}@berri.ai".format(uuid.uuid4())
|
|
|
|
|
new_user_info = await new_user(session=session, i=0, user_email=user_email)
|
|
|
|
|
new_member = {"role": "user", "user_email": user_email}
|
|
|
|
|
await add_member(
|
|
|
|
|
session=session, i=0, team_id=team_data["team_id"], members=[new_member]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
## check user info to confirm user is in team
|
|
|
|
|
updated_user_info = await get_user_info(
|
|
|
|
|
session=session, get_user=new_user_info["user_id"], call_user="sk-1234"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
print(updated_user_info)
|
|
|
|
|
|
|
|
|
|
## check if team in user table
|
|
|
|
|
is_team_in_list: bool = False
|
|
|
|
|
for team in updated_user_info["teams"]:
|
|
|
|
|
if team_data["team_id"] == team["team_id"]:
|
|
|
|
|
is_team_in_list = True
|
|
|
|
|
assert is_team_in_list
|
|
|
|
|
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_team_delete():
|
|
|
|
|
"""
|
|
|
|
|
- Create team
|
|
|
|
|
- Create key for team
|
|
|
|
|
- Check if key works
|
|
|
|
|
- Delete team
|
|
|
|
|
"""
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
|
## Create admin
|
|
|
|
|
admin_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=admin_user)
|
|
|
|
|
## Create normal user
|
|
|
|
|
normal_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=normal_user)
|
|
|
|
|
## Create team with 1 admin and 1 user
|
|
|
|
|
member_list = [
|
|
|
|
|
{"role": "admin", "user_id": admin_user},
|
|
|
|
|
{"role": "user", "user_id": normal_user},
|
|
|
|
|
]
|
|
|
|
|
team_data = await new_team(session=session, i=0, member_list=member_list)
|
2025-02-18 19:14:20 -08:00
|
|
|
|
|
|
|
|
## ASSERT USER MEMBERSHIP IS CREATED
|
|
|
|
|
user_info = await get_user_info(
|
|
|
|
|
session=session, get_user=normal_user, call_user="sk-1234"
|
|
|
|
|
)
|
|
|
|
|
assert len(user_info["teams"]) == 1
|
|
|
|
|
|
2024-02-21 16:53:12 -08:00
|
|
|
## Create key
|
|
|
|
|
key_gen = await generate_key(session=session, i=0, team_id=team_data["team_id"])
|
|
|
|
|
key = key_gen["key"]
|
|
|
|
|
## Test key
|
2025-03-06 22:04:36 -08:00
|
|
|
# response = await chat_completion(session=session, key=key)
|
2024-02-21 16:53:12 -08:00
|
|
|
## Delete team
|
|
|
|
|
await delete_team(session=session, i=0, team_id=team_data["team_id"])
|
2024-03-01 09:14:08 -08:00
|
|
|
|
2025-02-18 19:14:20 -08:00
|
|
|
## ASSERT USER MEMBERSHIP IS DELETED
|
|
|
|
|
user_info = await get_user_info(
|
|
|
|
|
session=session, get_user=normal_user, call_user="sk-1234"
|
|
|
|
|
)
|
|
|
|
|
assert len(user_info["teams"]) == 0
|
|
|
|
|
|
2025-03-06 22:04:36 -08:00
|
|
|
## ASSERT TEAM INFO NOW RETURNS A 404
|
|
|
|
|
with pytest.raises(openai.NotFoundError):
|
|
|
|
|
await get_team_info(
|
|
|
|
|
session=session, get_team=team_data["team_id"], call_key="sk-1234"
|
|
|
|
|
)
|
|
|
|
|
|
2024-03-01 09:14:08 -08:00
|
|
|
|
2024-06-08 16:27:14 -07:00
|
|
|
@pytest.mark.parametrize("dimension", ["user_id", "user_email"])
|
2024-03-01 09:14:08 -08:00
|
|
|
@pytest.mark.asyncio
|
2024-06-08 16:27:14 -07:00
|
|
|
async def test_member_delete(dimension):
|
2024-03-01 09:14:08 -08:00
|
|
|
"""
|
|
|
|
|
- Create team
|
|
|
|
|
- Add member
|
|
|
|
|
- Get team info (check if member in team)
|
|
|
|
|
- Delete member
|
|
|
|
|
- Get team info (check if member in team)
|
|
|
|
|
"""
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
|
# Create Team
|
|
|
|
|
## Create admin
|
|
|
|
|
admin_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=admin_user)
|
|
|
|
|
## Create normal user
|
|
|
|
|
normal_user = f"{uuid.uuid4()}"
|
2024-06-08 16:27:14 -07:00
|
|
|
normal_user_email = "{}@berri.ai".format(normal_user)
|
2024-03-01 09:14:08 -08:00
|
|
|
print(f"normal_user: {normal_user}")
|
2024-06-08 16:27:14 -07:00
|
|
|
await new_user(
|
|
|
|
|
session=session, i=0, user_id=normal_user, user_email=normal_user_email
|
|
|
|
|
)
|
2024-03-01 09:14:08 -08:00
|
|
|
## Create team with 1 admin and 1 user
|
|
|
|
|
member_list = [
|
|
|
|
|
{"role": "admin", "user_id": admin_user},
|
|
|
|
|
]
|
2024-06-08 16:27:14 -07:00
|
|
|
if dimension == "user_id":
|
|
|
|
|
member_list.append({"role": "user", "user_id": normal_user})
|
|
|
|
|
elif dimension == "user_email":
|
|
|
|
|
member_list.append({"role": "user", "user_email": normal_user_email})
|
2024-03-01 09:14:08 -08:00
|
|
|
team_data = await new_team(session=session, i=0, member_list=member_list)
|
2024-06-08 16:27:14 -07:00
|
|
|
|
|
|
|
|
user_in_team = False
|
2024-03-01 09:14:08 -08:00
|
|
|
for member in team_data["members_with_roles"]:
|
2024-06-08 16:27:14 -07:00
|
|
|
if dimension == "user_id" and member["user_id"] == normal_user:
|
|
|
|
|
user_in_team = True
|
|
|
|
|
elif (
|
|
|
|
|
dimension == "user_email" and member["user_email"] == normal_user_email
|
|
|
|
|
):
|
|
|
|
|
user_in_team = True
|
2024-03-01 09:14:08 -08:00
|
|
|
|
2024-06-08 16:27:14 -07:00
|
|
|
assert (
|
|
|
|
|
user_in_team is True
|
|
|
|
|
), "User not in team. Team list={}, User details - id={}, email={}. Dimension={}".format(
|
|
|
|
|
team_data["members_with_roles"], normal_user, normal_user_email, dimension
|
2024-03-01 09:14:08 -08:00
|
|
|
)
|
2024-06-08 16:27:14 -07:00
|
|
|
# Delete member
|
|
|
|
|
if dimension == "user_id":
|
|
|
|
|
updated_team_data = await delete_member(
|
|
|
|
|
session=session, i=0, team_id=team_data["team_id"], user_id=normal_user
|
|
|
|
|
)
|
|
|
|
|
elif dimension == "user_email":
|
|
|
|
|
updated_team_data = await delete_member(
|
|
|
|
|
session=session,
|
|
|
|
|
i=0,
|
|
|
|
|
team_id=team_data["team_id"],
|
|
|
|
|
user_email=normal_user_email,
|
|
|
|
|
)
|
2024-03-01 09:14:08 -08:00
|
|
|
print(f"updated_team_data: {updated_team_data}")
|
2024-06-08 16:27:14 -07:00
|
|
|
user_in_team = False
|
|
|
|
|
for member in team_data["members_with_roles"]:
|
|
|
|
|
if dimension == "user_id" and member["user_id"] == normal_user:
|
|
|
|
|
user_in_team = True
|
|
|
|
|
elif (
|
|
|
|
|
dimension == "user_email" and member["user_email"] == normal_user_email
|
|
|
|
|
):
|
|
|
|
|
user_in_team = True
|
2024-03-01 09:14:08 -08:00
|
|
|
|
2024-06-08 16:27:14 -07:00
|
|
|
assert user_in_team is True
|
2024-03-07 07:56:51 -08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_team_alias():
|
|
|
|
|
"""
|
|
|
|
|
- Create team w/ model alias
|
|
|
|
|
- Create key for team
|
|
|
|
|
- Check if key works
|
|
|
|
|
"""
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
|
## Create admin
|
|
|
|
|
admin_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=admin_user)
|
|
|
|
|
## Create normal user
|
|
|
|
|
normal_user = f"{uuid.uuid4()}"
|
|
|
|
|
await new_user(session=session, i=0, user_id=normal_user)
|
|
|
|
|
## Create team with 1 admin and 1 user
|
|
|
|
|
member_list = [
|
|
|
|
|
{"role": "admin", "user_id": admin_user},
|
|
|
|
|
{"role": "user", "user_id": normal_user},
|
|
|
|
|
]
|
|
|
|
|
team_data = await new_team(
|
|
|
|
|
session=session,
|
|
|
|
|
i=0,
|
|
|
|
|
member_list=member_list,
|
|
|
|
|
model_aliases={"cheap-model": "gpt-3.5-turbo"},
|
|
|
|
|
)
|
|
|
|
|
## Create key
|
|
|
|
|
key_gen = await generate_key(
|
|
|
|
|
session=session, i=0, team_id=team_data["team_id"], models=["gpt-3.5-turbo"]
|
|
|
|
|
)
|
|
|
|
|
key = key_gen["key"]
|
|
|
|
|
## Test key
|
|
|
|
|
response = await chat_completion(session=session, key=key, model="cheap-model")
|
2024-05-22 19:19:51 -07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
|
async def test_users_in_team_budget():
|
|
|
|
|
"""
|
|
|
|
|
- Create User
|
2025-12-06 22:28:23 -08:00
|
|
|
- Create Team with User
|
2024-05-22 19:19:51 -07:00
|
|
|
- Add User to team with budget = 0.0000001
|
|
|
|
|
- Make Call 1 -> pass
|
|
|
|
|
- Make Call 2 -> fail
|
|
|
|
|
"""
|
|
|
|
|
get_user = f"krrish_{time.time()}@berri.ai"
|
|
|
|
|
async with aiohttp.ClientSession() as session:
|
2025-12-20 12:30:30 -08:00
|
|
|
# IMPORTANT: Create team first, then create user with team_id.
|
|
|
|
|
# This order is critical for the test to work correctly:
|
|
|
|
|
# - When a user is created with team_id, the API key gets team_id set from the start
|
|
|
|
|
# - This ensures spend tracking and budget enforcement work correctly
|
|
|
|
|
# - If we create the user first (without team_id) and then add them to a team,
|
|
|
|
|
# the key's team_id remains None, breaking team budget tracking
|
|
|
|
|
# DO NOT change this order - it's testing the intended flow where keys are
|
|
|
|
|
# associated with teams at creation time.
|
|
|
|
|
team = await new_team(session, 0, user_id=None)
|
|
|
|
|
print(f"[DEBUG] Created team: {team['team_id']}")
|
|
|
|
|
print(f"[DEBUG] Full team data: {team}")
|
|
|
|
|
|
|
|
|
|
# Create user with team_id so the key is associated with the team from the start
|
2024-05-22 19:19:51 -07:00
|
|
|
key_gen = await new_user(
|
|
|
|
|
session,
|
|
|
|
|
0,
|
|
|
|
|
user_id=get_user,
|
|
|
|
|
budget=10,
|
|
|
|
|
budget_duration="5s",
|
|
|
|
|
models=["fake-openai-endpoint"],
|
2025-12-20 12:30:30 -08:00
|
|
|
team_id=team["team_id"],
|
2024-05-22 19:19:51 -07:00
|
|
|
)
|
|
|
|
|
key = key_gen["key"]
|
2025-12-20 12:30:30 -08:00
|
|
|
print(f"[DEBUG] Created user '{get_user}' with key: {key}")
|
|
|
|
|
print(f"[DEBUG] User budget: 10, budget_duration: 5s")
|
|
|
|
|
print(f"[DEBUG] Key team_id: {team['team_id']}")
|
|
|
|
|
|
|
|
|
|
# Check user info BEFORE updating member budget
|
|
|
|
|
user_info_before = await get_user_info(session, get_user, call_user="sk-1234")
|
|
|
|
|
print(f"[DEBUG] User info BEFORE update_member:")
|
|
|
|
|
print(f" - User budget: {user_info_before.get('max_budget')}")
|
|
|
|
|
print(f" - User spend: {user_info_before.get('spend')}")
|
|
|
|
|
if user_info_before.get("teams"):
|
|
|
|
|
for team_info in user_info_before["teams"]:
|
|
|
|
|
if team_info.get("team_id") == team["team_id"]:
|
|
|
|
|
print(f" - Team memberships: {team_info.get('team_memberships')}")
|
2024-05-22 19:19:51 -07:00
|
|
|
|
2025-02-01 11:23:00 -08:00
|
|
|
# update user to have budget = 0.0000001
|
2025-12-20 12:30:30 -08:00
|
|
|
update_result = await update_member(
|
2024-05-22 19:19:51 -07:00
|
|
|
session, 0, team_id=team["team_id"], user_id=get_user, max_budget=0.0000001
|
|
|
|
|
)
|
2025-12-20 12:30:30 -08:00
|
|
|
print(f"[DEBUG] Updated member budget to 0.0000001")
|
|
|
|
|
print(f"[DEBUG] Update result: {update_result}")
|
|
|
|
|
|
|
|
|
|
# Check user info AFTER updating member budget
|
|
|
|
|
user_info_after = await get_user_info(session, get_user, call_user="sk-1234")
|
|
|
|
|
print(f"[DEBUG] User info AFTER update_member:")
|
|
|
|
|
print(f" - User budget: {user_info_after.get('max_budget')}")
|
|
|
|
|
print(f" - User spend: {user_info_after.get('spend')}")
|
|
|
|
|
if user_info_after.get("teams"):
|
|
|
|
|
for team_info in user_info_after["teams"]:
|
|
|
|
|
if team_info.get("team_id") == team["team_id"]:
|
|
|
|
|
print(f" - Team: {team_info.get('team_id')}")
|
|
|
|
|
for membership in team_info.get('team_memberships', []):
|
|
|
|
|
print(f" - Membership: {membership}")
|
|
|
|
|
if 'litellm_budget_table' in membership:
|
|
|
|
|
budget_table = membership['litellm_budget_table']
|
|
|
|
|
print(f" - Max budget: {budget_table.get('max_budget')}")
|
|
|
|
|
print(f" - Current spend: {membership.get('spend', 0)}")
|
2024-05-22 19:19:51 -07:00
|
|
|
|
|
|
|
|
# Call 1
|
2025-12-20 12:30:30 -08:00
|
|
|
print("\n[DEBUG] ===== Making Call 1 =====")
|
2024-05-22 19:19:51 -07:00
|
|
|
result = await chat_completion(session, key, model="fake-openai-endpoint")
|
2025-12-20 12:30:30 -08:00
|
|
|
print(f"[DEBUG] Call 1 PASSED (expected)")
|
|
|
|
|
print(f"[DEBUG] Call 1 result: {result}")
|
|
|
|
|
# Extract cost from result if available
|
|
|
|
|
if isinstance(result, dict):
|
|
|
|
|
usage = result.get('usage', {})
|
|
|
|
|
print(f"[DEBUG] Call 1 usage: {usage}")
|
2024-05-22 19:19:51 -07:00
|
|
|
|
2025-12-18 13:16:15 -08:00
|
|
|
# Wait for spend to be committed to database before checking budget
|
|
|
|
|
# Spend updates are queued asynchronously and committed periodically (every minute),
|
|
|
|
|
# so we need to wait for the spend from Call 1 to be persisted
|
2025-12-20 12:30:30 -08:00
|
|
|
print("\n[DEBUG] ===== Waiting for spend to be committed =====")
|
2025-12-18 13:16:15 -08:00
|
|
|
print("Waiting for team member spend to be committed to database...")
|
2026-03-13 12:40:17 -07:00
|
|
|
print("Note: Spend updates are flushed periodically, this may take up to 90 seconds...")
|
2025-12-18 13:16:15 -08:00
|
|
|
spend_updated = await wait_for_team_member_spend_update(
|
2026-03-13 12:40:17 -07:00
|
|
|
session, get_user, team["team_id"], 0.0000001, max_wait=90
|
2025-12-18 13:16:15 -08:00
|
|
|
)
|
|
|
|
|
if not spend_updated:
|
2026-03-13 12:40:17 -07:00
|
|
|
pytest.fail(
|
|
|
|
|
"Team member spend was not updated within 90s. "
|
|
|
|
|
"The spend update queue may not have flushed, or the model may have 0 cost."
|
|
|
|
|
)
|
2024-05-22 19:19:51 -07:00
|
|
|
|
2025-12-20 12:30:30 -08:00
|
|
|
# Check user info BEFORE Call 2
|
|
|
|
|
user_info_before_call2 = await get_user_info(session, get_user, call_user="sk-1234")
|
|
|
|
|
print(f"\n[DEBUG] User info BEFORE Call 2:")
|
|
|
|
|
print(f" - User budget: {user_info_before_call2.get('max_budget')}")
|
|
|
|
|
print(f" - User spend: {user_info_before_call2.get('spend')}")
|
|
|
|
|
if user_info_before_call2.get("teams"):
|
|
|
|
|
for team_info in user_info_before_call2["teams"]:
|
|
|
|
|
if team_info.get("team_id") == team["team_id"]:
|
|
|
|
|
print(f" - Team: {team_info.get('team_id')}")
|
|
|
|
|
for membership in team_info.get('team_memberships', []):
|
|
|
|
|
if 'litellm_budget_table' in membership:
|
|
|
|
|
budget_table = membership['litellm_budget_table']
|
|
|
|
|
current_spend = membership.get('spend', 0)
|
|
|
|
|
max_budget = budget_table.get('max_budget')
|
|
|
|
|
print(f" - Max budget in team: {max_budget}")
|
|
|
|
|
print(f" - Current spend in team: {current_spend}")
|
|
|
|
|
print(f" - Budget remaining: {max_budget - current_spend}")
|
|
|
|
|
print(f" - Should fail?: {current_spend >= max_budget}")
|
|
|
|
|
|
2024-05-22 19:19:51 -07:00
|
|
|
# Call 2
|
2025-12-20 12:30:30 -08:00
|
|
|
print("\n[DEBUG] ===== Making Call 2 =====")
|
|
|
|
|
call2_failed = False
|
|
|
|
|
call2_error = None
|
|
|
|
|
call2_status = None
|
2024-05-22 19:19:51 -07:00
|
|
|
try:
|
2025-12-20 12:30:30 -08:00
|
|
|
# Capture the response to check status code
|
|
|
|
|
url = "http://localhost:4000/chat/completions"
|
|
|
|
|
headers = {
|
|
|
|
|
"Authorization": f"Bearer {key}",
|
|
|
|
|
"Content-Type": "application/json",
|
|
|
|
|
}
|
|
|
|
|
data = {
|
|
|
|
|
"model": "fake-openai-endpoint",
|
|
|
|
|
"messages": [
|
|
|
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
|
|
|
{"role": "user", "content": "Hello!"},
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
|
|
|
call2_status = response.status
|
|
|
|
|
response_text = await response.text()
|
|
|
|
|
print(f"[DEBUG] Call 2 status code: {call2_status}")
|
|
|
|
|
print(f"[DEBUG] Call 2 response: {response_text}")
|
|
|
|
|
|
|
|
|
|
if call2_status != 200:
|
|
|
|
|
call2_failed = True
|
|
|
|
|
call2_error = f"Status {call2_status}: {response_text}"
|
|
|
|
|
raise Exception(call2_error)
|
|
|
|
|
else:
|
|
|
|
|
# Call succeeded when it should have failed
|
|
|
|
|
print(f"[ERROR] Call 2 PASSED when it should have FAILED!")
|
|
|
|
|
print(f"[ERROR] Response was 200 OK")
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
if call2_failed:
|
|
|
|
|
print(f"[DEBUG] Call 2 FAILED (expected): {e}")
|
|
|
|
|
print(f"[DEBUG] Checking if error message indicates budget exceeded...")
|
|
|
|
|
else:
|
|
|
|
|
call2_error = str(e)
|
|
|
|
|
print(f"[DEBUG] Call 2 raised exception: {e}")
|
|
|
|
|
|
|
|
|
|
# Check user info AFTER Call 2
|
|
|
|
|
user_info_after_call2 = await get_user_info(session, get_user, call_user="sk-1234")
|
|
|
|
|
print(f"\n[DEBUG] User info AFTER Call 2:")
|
|
|
|
|
print(f" - User budget: {user_info_after_call2.get('max_budget')}")
|
|
|
|
|
print(f" - User spend: {user_info_after_call2.get('spend')}")
|
|
|
|
|
if user_info_after_call2.get("teams"):
|
|
|
|
|
for team_info in user_info_after_call2["teams"]:
|
|
|
|
|
if team_info.get("team_id") == team["team_id"]:
|
|
|
|
|
print(f" - Team: {team_info.get('team_id')}")
|
|
|
|
|
for membership in team_info.get('team_memberships', []):
|
|
|
|
|
if 'litellm_budget_table' in membership:
|
|
|
|
|
budget_table = membership['litellm_budget_table']
|
|
|
|
|
print(f" - Max budget: {budget_table.get('max_budget')}")
|
|
|
|
|
print(f" - Current spend: {membership.get('spend', 0)}")
|
|
|
|
|
|
|
|
|
|
# Assert Call 2 failed
|
|
|
|
|
if not call2_failed:
|
|
|
|
|
error_msg = (
|
|
|
|
|
f"\n[FAILURE] Call 2 should have failed but it passed!\n"
|
|
|
|
|
f"Expected: Budget enforcement to block the call\n"
|
|
|
|
|
f"Actual: Call returned status {call2_status}\n"
|
|
|
|
|
f"Team member budget: 0.0000001\n"
|
|
|
|
|
f"User budget: {user_info_before_call2.get('max_budget')}\n"
|
|
|
|
|
f"User spend before call: {user_info_before_call2.get('spend')}\n"
|
|
|
|
|
)
|
|
|
|
|
# Add team member info if available
|
|
|
|
|
if user_info_before_call2.get("teams"):
|
|
|
|
|
for team_info in user_info_before_call2["teams"]:
|
|
|
|
|
if team_info.get("team_id") == team["team_id"]:
|
|
|
|
|
for membership in team_info.get('team_memberships', []):
|
|
|
|
|
if 'litellm_budget_table' in membership:
|
|
|
|
|
error_msg += f"Team member spend before call: {membership.get('spend', 0)}\n"
|
|
|
|
|
error_msg += f"Team member max budget: {membership['litellm_budget_table'].get('max_budget')}\n"
|
|
|
|
|
pytest.fail(error_msg)
|
|
|
|
|
|
|
|
|
|
# Check the error message contains budget exceeded
|
|
|
|
|
if call2_error and "Budget has been exceeded" not in call2_error:
|
2024-05-22 19:19:51 -07:00
|
|
|
pytest.fail(
|
2025-12-20 12:30:30 -08:00
|
|
|
f"Call 2 failed but not with expected error message.\n"
|
|
|
|
|
f"Expected error to contain: 'Budget has been exceeded'\n"
|
|
|
|
|
f"Actual error: {call2_error}"
|
2024-05-22 19:19:51 -07:00
|
|
|
)
|
2025-12-20 12:30:30 -08:00
|
|
|
|
|
|
|
|
print("[DEBUG] Call 2 failed as expected with budget exceeded error")
|
2024-10-08 21:57:03 -07:00
|
|
|
|
|
|
|
|
## Check user info
|
|
|
|
|
user_info = await get_user_info(session, get_user, call_user="sk-1234")
|
|
|
|
|
|
|
|
|
|
assert (
|
|
|
|
|
user_info["teams"][0]["team_memberships"][0]["litellm_budget_table"][
|
|
|
|
|
"max_budget"
|
|
|
|
|
]
|
|
|
|
|
== 0.0000001
|
|
|
|
|
)
|
