From 5fcb4130a09f4ccdc1ee690448ffad3e636f36a7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 09:59:26 -0800 Subject: [PATCH 1/4] fix - store budget_duration for keys --- .../key_management_endpoints.py | 1 + .../test_key_generate_prisma.py | 19 ++++++++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ba6df9a2e..f7d1658e1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -476,6 +476,7 @@ def prepare_key_update_data( duration_s = duration_in_seconds(duration=budget_duration) key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) non_default_values["budget_reset_at"] = key_reset_at + non_default_values["budget_duration"] = budget_duration _metadata = existing_key_row.metadata or {} diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index e6f8ca541..26e3f0646 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -23,7 +23,7 @@ import os import sys import traceback import uuid -from datetime import datetime +from datetime import datetime, timezone from dotenv import load_dotenv from fastapi import Request @@ -1305,6 +1305,8 @@ def test_generate_and_update_key(prisma_client): data=UpdateKeyRequest( key=generated_key, models=["ada", "babbage", "curie", "davinci"], + budget_duration="1mo", + max_budget=100, ), ) @@ -1333,6 +1335,21 @@ def test_generate_and_update_key(prisma_client): } assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"] assert result["info"]["team_id"] == _team_2 + assert result["info"]["budget_duration"] == "1mo" + assert result["info"]["max_budget"] == 100 + + # budget_reset_at should be 30 days from now + assert result["info"]["budget_reset_at"] is not None + budget_reset_at = result["info"]["budget_reset_at"].replace( + tzinfo=timezone.utc + ) + current_time = datetime.now(timezone.utc) + + print( + "days between now and budget_reset_at", + (budget_reset_at - current_time).days, + ) + assert (budget_reset_at - current_time).days >= 29 # around 1 month # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) From 4f840b1679883544acadb961fa15c81f561b9bb5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 10:06:30 -0800 Subject: [PATCH 2/4] test_generate_and_update_key --- .../test_key_generate_prisma.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 26e3f0646..8d264ea75 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -1349,7 +1349,13 @@ def test_generate_and_update_key(prisma_client): "days between now and budget_reset_at", (budget_reset_at - current_time).days, ) - assert (budget_reset_at - current_time).days >= 29 # around 1 month + # assert budget_reset_at is 30 days from now + assert ( + abs( + (budget_reset_at - current_time).total_seconds() - 30 * 24 * 60 * 60 + ) + <= 10 + ) # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) @@ -2630,6 +2636,15 @@ async def test_create_update_team(prisma_client): _updated_info["budget_reset_at"], datetime.datetime ) + # budget_reset_at should be 2 days from now + budget_reset_at = _updated_info["budget_reset_at"].replace(tzinfo=timezone.utc) + current_time = datetime.datetime.now(timezone.utc) + + # assert budget_reset_at is 2 days from now + assert ( + abs((budget_reset_at - current_time).total_seconds() - 2 * 24 * 60 * 60) <= 10 + ) + # now hit team_info try: response = await team_info( From 12069c175dbada635e5103bcb183e59b321a11f4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 10:43:12 -0800 Subject: [PATCH 3/4] test_update_user_unit_test --- .../test_key_generate_prisma.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 8d264ea75..12122e80b 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -2788,6 +2788,58 @@ async def test_update_user_role(prisma_client): print("result from user auth with new key", result) +@pytest.mark.asyncio() +async def test_update_user_unit_test(prisma_client): + """ + Tests if we update user role, incorrect values are not stored in cache + -> create a user with role == INTERNAL_USER + -> access an Admin only route -> expect to fail + -> update user role to == PROXY_ADMIN + -> access an Admin only route -> expect to succeed + """ + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + key = await new_user( + data=NewUserRequest( + user_email="test@test.com", + ) + ) + + print(key) + + user_info = await user_update( + data=UpdateUserRequest( + user_id=key.user_id, + team_id="1234", + max_budget=100, + budget_duration="10d", + tpm_limit=100, + rpm_limit=100, + metadata={"very-new-metadata": "something"}, + ) + ) + + print("user_info", user_info) + assert user_info is not None + _user_info = user_info["data"].model_dump() + + assert _user_info["user_id"] == key.user_id + assert _user_info["team_id"] == "1234" + assert _user_info["max_budget"] == 100 + assert _user_info["budget_duration"] == "10d" + assert _user_info["tpm_limit"] == 100 + assert _user_info["rpm_limit"] == 100 + assert _user_info["metadata"] == {"very-new-metadata": "something"} + + # budget reset at should be 10 days from now + budget_reset_at = _user_info["budget_reset_at"].replace(tzinfo=timezone.utc) + current_time = datetime.now(timezone.utc) + assert ( + abs((budget_reset_at - current_time).total_seconds() - 10 * 24 * 60 * 60) <= 10 + ) + + @pytest.mark.asyncio() async def test_custom_api_key_header_name(prisma_client): """ """ From 052d9f62d90ca462c9944b8327e0b8df6375ce20 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 10:44:53 -0800 Subject: [PATCH 4/4] fix user update --- tests/proxy_unit_tests/test_key_generate_prisma.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 12122e80b..e34ed2b3e 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -2791,11 +2791,9 @@ async def test_update_user_role(prisma_client): @pytest.mark.asyncio() async def test_update_user_unit_test(prisma_client): """ - Tests if we update user role, incorrect values are not stored in cache - -> create a user with role == INTERNAL_USER - -> access an Admin only route -> expect to fail - -> update user role to == PROXY_ADMIN - -> access an Admin only route -> expect to succeed + Unit test for /user/update + + Ensure that params are updated for UpdateUserRequest """ setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")