litellm/tests/test_organizations.py
Krish Dholakia cf579fe644
Litellm stable pr 10 30 2024 (#6821)
* Update organization_endpoints.py to be able to list organizations (#6473)

* Update organization_endpoints.py to be able to list organizations

* Update test_organizations.py

* Update test_organizations.py

add test for list

* Update test_organizations.py

correct indentation

* Add unreleased Claude 3.5 Haiku models. (#6476)

---------

Co-authored-by: superpoussin22 <vincent.nadal@orange.fr>
Co-authored-by: David Manouchehri <david.manouchehri@ai.moda>
2024-11-20 05:03:42 +05:30

84 lines
2.6 KiB
Python

# What this tests ?
## Tests /organization endpoints.
import pytest
import asyncio
import aiohttp
import time, uuid
from openai import AsyncOpenAI
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 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.post(url, headers=headers) as response:
status = response.status
response_json = await response.json()
print(f"Response {i} (Status code: {status}):")
print(response_json)
print()
if status != 200:
raise Exception(f"Request {i} did not return a 200 status code: {status}")
return await response.json()
@pytest.mark.asyncio
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)
@pytest.mark.asyncio
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)
print(len(response_json))
if len(response_json)==0:
raise Exception(f"Return empty list of organization")