Spaces:
Configuration error
Configuration error
| # What this tests ? | |
| ## Tests /organization endpoints. | |
| import pytest | |
| import asyncio | |
| import aiohttp | |
| import time, uuid | |
| from openai import AsyncOpenAI | |
| async def new_user( | |
| session, | |
| i, | |
| user_id=None, | |
| budget=None, | |
| budget_duration=None, | |
| models=["azure-models"], | |
| team_id=None, | |
| user_email=None, | |
| ): | |
| url = "http://0.0.0.0:4000/user/new" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "models": models, | |
| "aliases": {"mistral-7b": "gpt-3.5-turbo"}, | |
| "duration": None, | |
| "max_budget": budget, | |
| "budget_duration": budget_duration, | |
| "user_email": user_email, | |
| } | |
| if user_id is not None: | |
| data["user_id"] = user_id | |
| if team_id is not None: | |
| data["team_id"] = team_id | |
| 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}, response: {response_text}" | |
| ) | |
| return await response.json() | |
| async def new_organization(session, i, organization_alias, max_budget=None): | |
| url = "http://0.0.0.0:4000/organization/new" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "organization_alias": organization_alias, | |
| "models": ["azure-models"], | |
| "max_budget": max_budget, | |
| } | |
| 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 add_member_to_org( | |
| session, i, organization_id, user_id, user_role="internal_user" | |
| ): | |
| url = "http://0.0.0.0:4000/organization/member_add" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "organization_id": organization_id, | |
| "member": { | |
| "user_id": user_id, | |
| "role": user_role, | |
| }, | |
| } | |
| 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 update_member_role( | |
| session, i, organization_id, user_id, user_role="internal_user" | |
| ): | |
| url = "http://0.0.0.0:4000/organization/member_update" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "organization_id": organization_id, | |
| "user_id": user_id, | |
| "role": user_role, | |
| } | |
| async with session.patch(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_member_from_org(session, i, organization_id, user_id): | |
| url = "http://0.0.0.0:4000/organization/member_delete" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = { | |
| "organization_id": organization_id, | |
| "user_id": user_id, | |
| } | |
| async with session.delete(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_organization(session, i, organization_id): | |
| url = "http://0.0.0.0:4000/organization/delete" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| data = {"organization_ids": [organization_id]} | |
| async with session.delete(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 list_organization(session, i): | |
| url = "http://0.0.0.0:4000/organization/list" | |
| headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} | |
| async with session.get(url, headers=headers) as response: | |
| status = response.status | |
| response_json = await response.json() | |
| print(f"Response {i} (Status code: {status}):") | |
| print() | |
| if status != 200: | |
| raise Exception(f"Request {i} did not return a 200 status code: {status}") | |
| # Assert that budget info is returned for each organization | |
| for org in response_json: | |
| assert "litellm_budget_table" in org, "Missing budget info in organization response" | |
| # Optionally also check that it's not null | |
| assert org["litellm_budget_table"] is not None, "Budget info is None" | |
| return response_json | |
| async def test_organization_new(): | |
| """ | |
| Make 20 parallel calls to /organization/new. Assert all worked. | |
| """ | |
| organization_alias = f"Organization: {uuid.uuid4()}" | |
| async with aiohttp.ClientSession() as session: | |
| tasks = [ | |
| new_organization( | |
| session=session, i=0, organization_alias=organization_alias | |
| ) | |
| for i in range(1, 20) | |
| ] | |
| await asyncio.gather(*tasks) | |
| async def test_organization_list(): | |
| """ | |
| create 2 new Organizations | |
| check if the Organization list is not empty | |
| """ | |
| organization_alias = f"Organization: {uuid.uuid4()}" | |
| async with aiohttp.ClientSession() as session: | |
| tasks = [ | |
| new_organization( | |
| session=session, i=0, organization_alias=organization_alias | |
| ) | |
| for i in range(1, 2) | |
| ] | |
| await asyncio.gather(*tasks) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| if len(response_json) == 0: | |
| raise Exception("Return empty list of organization") | |
| async def test_organization_delete(): | |
| """ | |
| create a new organization | |
| delete the organization | |
| check if the Organization list is set | |
| """ | |
| organization_alias = f"Organization: {uuid.uuid4()}" | |
| async with aiohttp.ClientSession() as session: | |
| tasks = [ | |
| new_organization( | |
| session=session, i=0, organization_alias=organization_alias | |
| ) | |
| ] | |
| await asyncio.gather(*tasks) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| organization_id = response_json[0]["organization_id"] | |
| await delete_organization(session, i=0, organization_id=organization_id) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| async def test_organization_member_flow(): | |
| """ | |
| create a new organization | |
| add a new member to the organization | |
| check if the member is added to the organization | |
| update the member's role in the organization | |
| delete the member from the organization | |
| check if the member is deleted from the organization | |
| """ | |
| organization_alias = f"Organization: {uuid.uuid4()}" | |
| async with aiohttp.ClientSession() as session: | |
| response_json = await new_organization( | |
| session=session, i=0, organization_alias=organization_alias | |
| ) | |
| organization_id = response_json["organization_id"] | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| new_user_response_json = await new_user( | |
| session=session, i=0, user_email=f"test_user_{uuid.uuid4()}@example.com" | |
| ) | |
| user_id = new_user_response_json["user_id"] | |
| await add_member_to_org( | |
| session, i=0, organization_id=organization_id, user_id=user_id | |
| ) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| for orgs in response_json: | |
| tmp_organization_id = orgs["organization_id"] | |
| if ( | |
| tmp_organization_id is not None | |
| and tmp_organization_id == organization_id | |
| ): | |
| user_id = orgs["members"][0]["user_id"] | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| await update_member_role( | |
| session, | |
| i=0, | |
| organization_id=organization_id, | |
| user_id=user_id, | |
| user_role="org_admin", | |
| ) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |
| await delete_member_from_org( | |
| session, i=0, organization_id=organization_id, user_id=user_id | |
| ) | |
| response_json = await list_organization(session, i=0) | |
| print(len(response_json)) | |