test(test_organizations.py): add testing for /organization/new endpoint

litellm/proxy/schema.prisma

@@ -21,12 +21,12 @@ model LiteLLM_BudgetTable {
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
   updated_by String
-  organization LiteLLM_OrganizationTable[]
+  organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget
 }
 
 model LiteLLM_OrganizationTable {
 		organization_id String @id @default(uuid())
-    organization_alias  String? 
+    organization_alias  String
     budget_id String
     metadata  Json  @default("{}")
     models     String[]

schema.prisma

@@ -21,11 +21,12 @@ model LiteLLM_BudgetTable {
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
   updated_by String
+  organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget
 }
 
 model LiteLLM_OrganizationTable {
 		organization_id String @id @default(uuid())
-    organization_alias  String? 
+    organization_alias  String
     budget_id String
     metadata  Json  @default("{}")
     models     String[]
@@ -35,7 +36,8 @@ model LiteLLM_OrganizationTable {
     created_by String
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     updated_by String
-    litellm_budget_table LiteLLM_BudgetTable   @relation(fields: [budget_id], references: [budget_id])
+    litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
+    teams LiteLLM_TeamTable[] 
 }
 
 // Assign prod keys to groups, not individuals 
@@ -59,7 +61,7 @@ model LiteLLM_TeamTable {
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
-    litellm_organization_table LiteLLM_OrganizationTable   @relation(fields: [organization_id], references: [organization_id])
+    litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])
 }
 
 // Track spend, rate limit, budget Users

tests/test_organizations.py

@@ -0,0 +1,46 @@
+# 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()
+
+
+@pytest.mark.asyncio
+async def test_organization_new():
+    """
+    Make 20 parallel calls to /user/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)