phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(proxy_server.py): fix team member delete endpoint

Browse source 
fixes issue in for loop, removing member from team
This commit is contained in:
Krrish Dholakia 2024-06-08 16:27:14 -07:00
parent d88b63264b
commit 5c4abcbdb4
2 changed files with 80 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

67
tests/test_team.py
View file

@ -15,6 +15,7 @@ async def new_user(
budget_duration=None,
models=["azure-models"],
team_id=None,
user_email=None,
):
url = "http://0.0.0.0:4000/user/new"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
@ -24,6 +25,7 @@ async def new_user(
"duration": None,
"max_budget": budget,
"budget_duration": budget_duration,
"user_email": user_email,
}
if user_id is not None:
@ -74,10 +76,14 @@ async def add_member(
return await response.json()
async def delete_member(session, i, team_id, user_id):
async def delete_member(session, i, team_id, user_id=None, user_email=None):
url = "http://0.0.0.0:4000/team/member_delete"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {"team_id": team_id, "user_id": user_id}
data = {"team_id": team_id}
if user_id is not None:
data["user_id"] = user_id
elif user_email is not None:
data["user_email"] = user_email
async with session.post(url, headers=headers, json=data) as response:
status = response.status
@ -415,8 +421,9 @@ async def test_team_delete():
await delete_team(session=session, i=0, team_id=team_data["team_id"])
@pytest.mark.parametrize("dimension", ["user_id", "user_email"])
@pytest.mark.asyncio
async def test_member_delete():
async def test_member_delete(dimension):
"""
- Create team
- Add member
@ -431,30 +438,58 @@ async def test_member_delete():
await new_user(session=session, i=0, user_id=admin_user)
## Create normal user
normal_user = f"{uuid.uuid4()}"
normal_user_email = "{}@berri.ai".format(normal_user)
print(f"normal_user: {normal_user}")
await new_user(session=session, i=0, user_id=normal_user)
await new_user(
session=session, i=0, user_id=normal_user, user_email=normal_user_email
)
## Create team with 1 admin and 1 user
member_list = [
{"role": "admin", "user_id": admin_user},
{"role": "user", "user_id": normal_user},
]
if dimension == "user_id":
member_list.append({"role": "user", "user_id": normal_user})
elif dimension == "user_email":
member_list.append({"role": "user", "user_email": normal_user_email})
team_data = await new_team(session=session, i=0, member_list=member_list)
print(f"team_data: {team_data}")
member_id_list = []
user_in_team = False
for member in team_data["members_with_roles"]:
member_id_list.append(member["user_id"])
if dimension == "user_id" and member["user_id"] == normal_user:
user_in_team = True
elif (
dimension == "user_email" and member["user_email"] == normal_user_email
):
user_in_team = True
assert normal_user in member_id_list
# Delete member
updated_team_data = await delete_member(
session=session, i=0, team_id=team_data["team_id"], user_id=normal_user
assert (
user_in_team is True
), "User not in team. Team list={}, User details - id={}, email={}. Dimension={}".format(
team_data["members_with_roles"], normal_user, normal_user_email, dimension
)
# Delete member
if dimension == "user_id":
updated_team_data = await delete_member(
session=session, i=0, team_id=team_data["team_id"], user_id=normal_user
)
elif dimension == "user_email":
updated_team_data = await delete_member(
session=session,
i=0,
team_id=team_data["team_id"],
user_email=normal_user_email,
)
print(f"updated_team_data: {updated_team_data}")
member_id_list = []
for member in updated_team_data["members_with_roles"]:
member_id_list.append(member["user_id"])
user_in_team = False
for member in team_data["members_with_roles"]:
if dimension == "user_id" and member["user_id"] == normal_user:
user_in_team = True
elif (
dimension == "user_email" and member["user_email"] == normal_user_email
):
user_in_team = True
assert normal_user not in member_id_list
assert user_in_team is True
@pytest.mark.asyncio