litellm

1
# What this tests ?
2
## Tests /organization endpoints.
3
import pytest
4
import asyncio
5
import aiohttp
6
import time, uuid
7
from openai import AsyncOpenAI
8

9

10
async def new_organization(session, i, organization_alias, max_budget=None):
11
    url = "http://0.0.0.0:4000/organization/new"
12
    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
13
    data = {
14
        "organization_alias": organization_alias,
15
        "models": ["azure-models"],
16
        "max_budget": max_budget,
17
    }
18

19
    async with session.post(url, headers=headers, json=data) as response:
20
        status = response.status
21
        response_text = await response.text()
22

23
        print(f"Response {i} (Status code: {status}):")
24
        print(response_text)
25
        print()
26

27
        if status != 200:
28
            raise Exception(f"Request {i} did not return a 200 status code: {status}")
29

30
        return await response.json()
31

32

33
@pytest.mark.asyncio
34
async def test_organization_new():
35
    """
36
    Make 20 parallel calls to /user/new. Assert all worked.
37
    """
38
    organization_alias = f"Organization: {uuid.uuid4()}"
39
    async with aiohttp.ClientSession() as session:
40
        tasks = [
41
            new_organization(
42
                session=session, i=0, organization_alias=organization_alias
43
            )
44
            for i in range(1, 20)
45
        ]
46
        await asyncio.gather(*tasks)
47