From d55237442bb2d50518c22647fc0116e1350ef00f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 20 Apr 2024 11:09:34 -0700 Subject: [PATCH] (ci/cd) testing with team_id and /user/new --- litellm/proxy/proxy_server.py | 7 ++++ litellm/tests/test_key_generate_prisma.py | 48 ++++++++++++++++++++--- 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 14ff78200b..2aab7e4532 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6532,6 +6532,13 @@ async def team_member_add( existing_team_row = await prisma_client.get_data( # type: ignore team_id=data.team_id, table_name="team", query_type="find_unique" ) + if existing_team_row is None: + raise HTTPException( + status_code=404, + detail={ + "error": f"Team not found for team_id={getattr(data, 'team_id', None)}" + }, + ) new_member = data.member diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index fdb7649d52..a90c13803f 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -120,6 +120,15 @@ async def test_new_user_response(prisma_client): await litellm.proxy.proxy_server.prisma_client.connect() from litellm.proxy.proxy_server import user_api_key_cache + await new_team( + NewTeamRequest( + team_id="ishaan-special-team", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + _response = await new_user( data=NewUserRequest( models=["azure-gpt-3.5"], @@ -999,10 +1008,32 @@ def test_generate_and_update_key(prisma_client): async def test(): await litellm.proxy.proxy_server.prisma_client.connect() + + # create team "litellm-core-infra@gmail.com"" + print("creating team litellm-core-infra@gmail.com") + await new_team( + NewTeamRequest( + team_id="litellm-core-infra@gmail.com", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + + await new_team( + NewTeamRequest( + team_id="ishaan-special-team", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + request = NewUserRequest( - metadata={"team": "litellm-team3", "project": "litellm-project3"}, + metadata={"project": "litellm-project3"}, team_id="litellm-core-infra@gmail.com", ) + key = await new_user(request) print(key) @@ -1015,7 +1046,6 @@ def test_generate_and_update_key(prisma_client): print("\n info for key=", result["info"]) assert result["info"]["max_parallel_requests"] == None assert result["info"]["metadata"] == { - "team": "litellm-team3", "project": "litellm-project3", } assert result["info"]["team_id"] == "litellm-core-infra@gmail.com" @@ -1037,7 +1067,7 @@ def test_generate_and_update_key(prisma_client): # update the team id response2 = await update_key_fn( request=Request, - data=UpdateKeyRequest(key=generated_key, team_id="ishaan"), + data=UpdateKeyRequest(key=generated_key, team_id="ishaan-special-team"), ) print("response2=", response2) @@ -1048,11 +1078,10 @@ def test_generate_and_update_key(prisma_client): print("\n info for key=", result["info"]) assert result["info"]["max_parallel_requests"] == None assert result["info"]["metadata"] == { - "team": "litellm-team3", "project": "litellm-project3", } assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"] - assert result["info"]["team_id"] == "ishaan" + assert result["info"]["team_id"] == "ishaan-special-team" # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) @@ -1941,6 +1970,15 @@ async def test_master_key_hashing(prisma_client): await litellm.proxy.proxy_server.prisma_client.connect() from litellm.proxy.proxy_server import user_api_key_cache + await new_team( + NewTeamRequest( + team_id="ishaan-special-team", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role="proxy_admin", api_key="sk-1234", user_id="1234" + ), + ) + _response = await new_user( data=NewUserRequest( models=["azure-gpt-3.5"],