From 61bc6fae636db56f6d787a8f83c6d96d6a10cf2f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 25 May 2024 15:20:49 -0700
Subject: [PATCH] test - new_team

---
 litellm/proxy/_types.py                   |  1 -
 litellm/proxy/proxy_server.py             | 16 +++-----
 litellm/tests/test_key_generate_prisma.py | 48 +++++++++++++++++++++++
 3 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 7be1eb03f6..11e358a2f8 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -565,7 +565,6 @@ class TeamBase(LiteLLMBase):
     # Budget fields
     max_budget: Optional[float] = None
     budget_duration: Optional[str] = None
-    budget_reset_at: Optional[datetime] = None
 
     models: list = []
     blocked: bool = False
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 136d379555..6c6a01f45b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -8034,20 +8034,16 @@ async def new_team(
         _model_id = model_dict.id
 
     ## ADD TO TEAM TABLE
-    # Check budget_duration and budget_reset_at
-    if data.budget_duration is not None:
-        duration_s = _duration_in_seconds(duration=data.budget_duration)
-        reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
-
-        # set the budget duration and budget_reset_at in DB
-        data.budget_duration = duration_s
-        data.budget_reset_at = reset_at
-
     complete_team_data = LiteLLM_TeamTable(
         **data.json(),
-        model_id=_model_id,
     )
 
+    # If budget_duration is set, set `budget_reset_at`
+    if complete_team_data.budget_duration is not None:
+        duration_s = _duration_in_seconds(duration=complete_team_data.budget_duration)
+        reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
+        complete_team_data.budget_reset_at = reset_at
+
     team_row = await prisma_client.insert_data(
         data=complete_team_data.json(exclude_none=True), table_name="team"
     )
diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py
index c1521e856c..3619ce8f49 100644
--- a/litellm/tests/test_key_generate_prisma.py
+++ b/litellm/tests/test_key_generate_prisma.py
@@ -2137,3 +2137,51 @@ async def test_reset_spend_authentication(prisma_client):
             "Tried to access route=/global/spend/reset, which is only for MASTER KEY"
             in e.message
         )
+
+
+@pytest.mark.asyncio()
+async def test_create_update_team(prisma_client):
+    """
+    - Set max_budget, budget_duration, max_budget, tpm_limit, rpm_limit
+    - Assert response has correct values
+
+    - Update max_budget, budget_duration, max_budget, tpm_limit, rpm_limit
+    - Assert response has correct values
+
+    - Call team_info and assert response has correct values
+    """
+    print("prisma client=", prisma_client)
+
+    master_key = "sk-1234"
+
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", master_key)
+    import datetime
+
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    from litellm.proxy.proxy_server import user_api_key_cache
+
+    _team_id = "test-team_{}".format(uuid.uuid4())
+    response = await new_team(
+        NewTeamRequest(
+            team_id=_team_id,
+            max_budget=20,
+            budget_duration="30d",
+            tpm_limit=20,
+            rpm_limit=20,
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+        ),
+    )
+
+    print("RESPONSE from new_team", response)
+
+    assert response["team_id"] == _team_id
+    assert response["max_budget"] == 20
+    assert response["tpm_limit"] == 20
+    assert response["rpm_limit"] == 20
+    assert response["budget_duration"] == "30d"
+    assert response["budget_reset_at"] is not None and isinstance(
+        response["budget_reset_at"], datetime.datetime
+    )