forked from BerriAI/litellm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_team_members.py
312 lines (248 loc) · 10.6 KB
/
test_team_members.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
import pytest
import requests
import time
from typing import Dict, List
import logging
import uuid
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class TeamAPI:
def __init__(self, base_url: str, auth_token: str):
self.base_url = base_url
self.headers = {
"Authorization": f"Bearer {auth_token}",
"Content-Type": "application/json",
}
def create_team(self, team_alias: str, models: List[str] = None) -> Dict:
"""Create a new team"""
# Generate a unique team_id using uuid
team_id = f"test_team_{uuid.uuid4().hex[:8]}"
data = {
"team_id": team_id,
"team_alias": team_alias,
"models": models or ["o3-mini"],
}
response = requests.post(
f"{self.base_url}/team/new", headers=self.headers, json=data
)
response.raise_for_status()
logger.info(f"Created new team: {team_id}")
return response.json(), team_id
def get_team_info(self, team_id: str) -> Dict:
"""Get current team information"""
response = requests.get(
f"{self.base_url}/team/info",
headers=self.headers,
params={"team_id": team_id},
)
response.raise_for_status()
return response.json()
def add_team_member(self, team_id: str, user_email: str, role: str) -> Dict:
"""Add a single team member"""
data = {"team_id": team_id, "member": [{"role": role, "user_id": user_email}]}
response = requests.post(
f"{self.base_url}/team/member_add", headers=self.headers, json=data
)
response.raise_for_status()
return response.json()
def delete_team_member(self, team_id: str, user_id: str) -> Dict:
"""Delete a team member
Args:
team_id (str): ID of the team
user_id (str): User ID to remove from team
Returns:
Dict: Response from the API
"""
data = {"team_id": team_id, "user_id": user_id}
response = requests.post(
f"{self.base_url}/team/member_delete", headers=self.headers, json=data
)
response.raise_for_status()
return response.json()
@pytest.fixture
def api_client():
"""Fixture for TeamAPI client"""
base_url = "http://localhost:4000"
auth_token = "sk-1234" # Replace with your token
return TeamAPI(base_url, auth_token)
@pytest.fixture
def new_team(api_client):
"""Fixture that creates a new team for each test"""
team_alias = f"Test Team {uuid.uuid4().hex[:6]}"
team_response, team_id = api_client.create_team(team_alias)
logger.info(f"Created test team: {team_id} ({team_alias})")
return team_id
def verify_member_in_team(team_info: Dict, user_email: str) -> bool:
"""Verify if a member exists in team"""
return any(
member["user_id"] == user_email
for member in team_info["team_info"]["members_with_roles"]
)
def test_team_creation(api_client):
"""Test team creation"""
team_alias = f"Test Team {uuid.uuid4().hex[:6]}"
team_response, team_id = api_client.create_team(team_alias)
# Verify team was created
team_info = api_client.get_team_info(team_id)
assert team_info["team_id"] == team_id
assert team_info["team_info"]["team_alias"] == team_alias
assert "o3-mini" in team_info["team_info"]["models"]
def test_add_single_member(api_client, new_team):
"""Test adding a single member to a new team"""
# Get initial team info
initial_info = api_client.get_team_info(new_team)
initial_size = len(initial_info["team_info"]["members_with_roles"])
# Add new member
test_email = f"pytest_user_{uuid.uuid4().hex[:6]}@mycompany.com"
api_client.add_team_member(new_team, test_email, "user")
# Allow time for system to process
time.sleep(1)
# Verify addition
updated_info = api_client.get_team_info(new_team)
updated_size = len(updated_info["team_info"]["members_with_roles"])
# Assertions
assert verify_member_in_team(
updated_info, test_email
), f"Member {test_email} not found in team"
assert (
updated_size == initial_size + 1
), f"Team size did not increase by 1 (was {initial_size}, now {updated_size})"
def test_add_multiple_members(api_client, new_team):
"""Test adding multiple members to a new team"""
# Get initial team size
initial_info = api_client.get_team_info(new_team)
initial_size = len(initial_info["team_info"]["members_with_roles"])
# Add 10 members
added_emails = []
for i in range(10):
email = f"pytest_user_{uuid.uuid4().hex[:6]}@mycompany.com"
added_emails.append(email)
logger.info(f"Adding member {i+1}/10: {email}")
api_client.add_team_member(new_team, email, "user")
# Allow time for system to process
time.sleep(1)
# Verify after each addition
current_info = api_client.get_team_info(new_team)
current_size = len(current_info["team_info"]["members_with_roles"])
# Assertions for each addition
assert verify_member_in_team(
current_info, email
), f"Member {email} not found in team"
assert (
current_size == initial_size + i + 1
), f"Team size incorrect after adding {email}"
# Final verification
final_info = api_client.get_team_info(new_team)
final_size = len(final_info["team_info"]["members_with_roles"])
# Final assertions
assert (
final_size == initial_size + 10
), f"Final team size incorrect (expected {initial_size + 10}, got {final_size})"
for email in added_emails:
assert verify_member_in_team(
final_info, email
), f"Member {email} not found in final team check"
def test_team_info_structure(api_client, new_team):
"""Test the structure of team info response"""
team_info = api_client.get_team_info(new_team)
# Verify required fields exist
assert "team_id" in team_info
assert "team_info" in team_info
assert "members_with_roles" in team_info["team_info"]
assert "models" in team_info["team_info"]
# Verify member structure
if team_info["team_info"]["members_with_roles"]:
member = team_info["team_info"]["members_with_roles"][0]
assert "user_id" in member
assert "role" in member
def test_error_handling(api_client):
"""Test error handling for invalid team ID"""
with pytest.raises(requests.exceptions.HTTPError):
api_client.get_team_info("invalid-team-id")
def test_duplicate_user_addition(api_client, new_team):
"""Test that adding the same user twice is handled appropriately"""
# Add user first time
test_email = f"pytest_user_{uuid.uuid4().hex[:6]}@mycompany.com"
initial_response = api_client.add_team_member(new_team, test_email, "user")
# Allow time for system to process
time.sleep(1)
# Get team info after first addition
team_info_after_first = api_client.get_team_info(new_team)
size_after_first = len(team_info_after_first["team_info"]["members_with_roles"])
logger.info(f"First addition completed. Team size: {size_after_first}")
# Attempt to add same user again
with pytest.raises(requests.exceptions.HTTPError):
api_client.add_team_member(new_team, test_email, "user")
# Allow time for system to process
time.sleep(1)
# Get team info after second addition attempt
team_info_after_second = api_client.get_team_info(new_team)
size_after_second = len(team_info_after_second["team_info"]["members_with_roles"])
# Verify team size didn't change
assert (
size_after_second == size_after_first
), f"Team size changed after duplicate addition (was {size_after_first}, now {size_after_second})"
# Verify user appears exactly once
user_count = sum(
1
for member in team_info_after_second["team_info"]["members_with_roles"]
if member["user_id"] == test_email
)
assert user_count == 1, f"User appears {user_count} times in team (expected 1)"
logger.info(f"Duplicate addition attempted. Final team size: {size_after_second}")
logger.info(f"Number of times user appears in team: {user_count}")
def test_member_deletion(api_client, new_team):
"""Test that member deletion works correctly and removes all instances of a user"""
# Add a test user
user_id = f"pytest_user_{uuid.uuid4().hex[:6]}"
api_client.add_team_member(new_team, user_id, "user")
time.sleep(1)
# Verify user was added
team_info_before = api_client.get_team_info(new_team)
assert verify_member_in_team(
team_info_before, user_id
), "User was not added successfully"
initial_size = len(team_info_before["team_info"]["members_with_roles"])
# Attempt to delete the same user multiple times (5 times)
for i in range(5):
logger.info(f"Attempting deletion {i+1}/5")
if i == 0:
# First deletion should succeed
api_client.delete_team_member(new_team, user_id)
time.sleep(1)
else:
# Subsequent deletions should raise an error
try:
api_client.delete_team_member(new_team, user_id)
pytest.fail("Expected HTTPError for duplicate deletion")
except requests.exceptions.HTTPError as e:
logger.info(
f"Expected error received on deletion attempt {i+1}: {str(e)}"
)
# Verify final state
final_info = api_client.get_team_info(new_team)
final_size = len(final_info["team_info"]["members_with_roles"])
# Verify user is completely removed
assert not verify_member_in_team(
final_info, user_id
), "User still exists in team after deletion"
# Verify only one member was removed
assert (
final_size == initial_size - 1
), f"Team size changed unexpectedly (was {initial_size}, now {final_size})"
def test_delete_nonexistent_member(api_client, new_team):
"""Test that attempting to delete a nonexistent member raises appropriate error"""
nonexistent_user = f"nonexistent_{uuid.uuid4().hex[:6]}"
# Verify user doesn't exist first
team_info = api_client.get_team_info(new_team)
assert not verify_member_in_team(
team_info, nonexistent_user
), "Test setup error: nonexistent user somehow exists"
# Attempt to delete nonexistent user
try:
api_client.delete_team_member(new_team, nonexistent_user)
pytest.fail("Expected HTTPError for deleting nonexistent user")
except requests.exceptions.HTTPError as e:
logger.info(f"Expected error received: {str(e)}")
assert e.response.status_code == 400