Spaces:
Configuration error
Configuration error
File size: 8,115 Bytes
447ebeb |
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 |
import pytest
import asyncio
import aiohttp
import json
from httpx import AsyncClient
from typing import Any, Optional
# =====================================================================
# NEW HELPER FUNCTIONS FOR TEAM BLOCKING TESTS
# =====================================================================
async def generate_team_key(
session,
team_id: str,
max_budget: Optional[float] = None,
):
"""Helper function to generate a key for a specific team"""
url = "http://0.0.0.0:4000/key/generate"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data: dict[str, Any] = {"team_id": team_id}
if max_budget is not None:
data["max_budget"] = max_budget
async with session.post(url, headers=headers, json=data) as response:
return await response.json()
async def update_team_block_status(session, team_id: str, blocked: bool, port: int):
"""Helper to update a team's 'blocked' status on a given instance port."""
url = f"http://0.0.0.0:{port}/team/update"
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
data = {"team_id": team_id, "blocked": blocked}
async with session.post(url, headers=headers, json=data) as response:
return await response.json()
async def get_team_info(session, team_id: str, port: int):
"""Helper to retrieve team info from a specific instance port."""
url = f"http://0.0.0.0:{port}/team/info"
headers = {"Authorization": "Bearer sk-1234"}
async with session.get(
url, headers=headers, params={"team_id": team_id}
) as response:
data = await response.json()
return data["team_info"]
async def chat_completion_on_port(
session, key: str, model: str, port: int, prompt: Optional[str] = None
):
"""
Helper function to make a chat completion request on a specified instance port.
Accepts an optional prompt string.
"""
from openai import AsyncOpenAI
import uuid
if prompt is None:
prompt = f"Say hello! {uuid.uuid4()}" * 100
client = AsyncOpenAI(api_key=key, base_url=f"http://0.0.0.0:{port}/v1")
response = await client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
)
return response
# =====================================================================
# NEW END‑TO‑END TEST FOR TEAM BLOCKING ACROSS MULTI‑INSTANCE SETUP
# =====================================================================
@pytest.mark.asyncio
async def test_team_blocking_behavior_multi_instance():
"""
Test team blocking scenario across multi-instance setup:
1. Create a new team on port 4000.
2. Verify (via team/info on port 4001) that the team is not blocked.
3. Create a key for that team.
4. Make a chat completion request (via instance on port 4000) and verify that it works.
6. Update the team to set 'blocked': True via the update endpoint on port 4001.
--- Sleep for 61 seconds --- the in-memory team obj ttl is 60 seconds
7. Verify (via team/info on port 4000) that the team is now blocked.
8. Make a chat completion request (using instance on port 4000) with a new prompt; expect it to be blocked.
9. Repeat the chat completion request with another new prompt; expect it to be blocked.
10. Confirm via team/info endpoints on both ports that the team remains blocked.
"""
async with aiohttp.ClientSession() as session:
headers = {
"Authorization": "Bearer sk-1234",
"Content-Type": "application/json",
}
# 1. Create a new team on instance (port 4000)
url_new_team = "http://0.0.0.0:4000/team/new"
team_data = {}
async with session.post(
url_new_team, headers=headers, json=team_data
) as response:
assert response.status == 200, "Failed to create team"
team_resp = await response.json()
team_id = team_resp["team_id"]
# 2. Verify via team/info on port 4001 that team is not blocked.
team_info_4001 = await get_team_info(session, team_id, port=4001)
assert "blocked" in team_info_4001, "Team info missing 'blocked' field"
assert (
team_info_4001["blocked"] is False
), "Team should not be blocked initially"
# 3. Create a key for the team using the existing helper.
key_gen = await generate_team_key(session=session, team_id=team_id)
key = key_gen["key"]
# 4. Make a chat completion request on port 4000 and verify it works.
response = await chat_completion_on_port(
session,
key=key,
model="fake-openai-endpoint",
port=4000,
prompt="Non-cached prompt 1",
)
assert (
response is not None
), "Chat completion should succeed when team is not blocked"
# 5. Update the team to set 'blocked': True on instance port 4001.
await update_team_block_status(session, team_id, blocked=True, port=4001)
print("sleeping for 61 seconds")
await asyncio.sleep(61)
# 6. Verify via team/info on port 4000 that the team is blocked.
team_info_4000 = await get_team_info(session, team_id, port=4000)
assert "blocked" in team_info_4000, "Team info missing 'blocked' field"
print(
"Team info on port 4000: ",
json.dumps(team_info_4000, indent=4, default=str),
)
assert team_info_4000["blocked"] is True, "Team should be blocked after update"
# 7. Verify via team/info on port 4001 that the team is blocked.
team_info_4001 = await get_team_info(session, team_id, port=4001)
assert "blocked" in team_info_4001, "Team info missing 'blocked' field"
assert team_info_4001["blocked"] is True, "Team should be blocked after update"
# 8. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked.
with pytest.raises(Exception) as excinfo:
await chat_completion_on_port(
session,
key=key,
model="fake-openai-endpoint",
port=4001,
prompt="Non-cached prompt 2",
)
error_msg = str(excinfo.value)
assert (
"blocked" in error_msg.lower()
), f"Expected error indicating team blocked, got: {error_msg}"
# 9. Make a chat completion request on port 4000 with a new prompt; expect it to be blocked.
with pytest.raises(Exception) as excinfo:
await chat_completion_on_port(
session,
key=key,
model="fake-openai-endpoint",
port=4000,
prompt="Non-cached prompt 2",
)
error_msg = str(excinfo.value)
assert (
"blocked" in error_msg.lower()
), f"Expected error indicating team blocked, got: {error_msg}"
# 9. Repeat the chat completion request with another new prompt; expect it to be blocked.
with pytest.raises(Exception) as excinfo_second:
await chat_completion_on_port(
session,
key=key,
model="fake-openai-endpoint",
port=4000,
prompt="Non-cached prompt 3",
)
error_msg_second = str(excinfo_second.value)
assert (
"blocked" in error_msg_second.lower()
), f"Expected error indicating team blocked, got: {error_msg_second}"
# 10. Final verification: check team info on both ports indicates the team is blocked.
final_team_info_4000 = await get_team_info(session, team_id, port=4000)
final_team_info_4001 = await get_team_info(session, team_id, port=4001)
assert (
final_team_info_4000.get("blocked") is True
), "Team on port 4000 should be blocked"
assert (
final_team_info_4001.get("blocked") is True
), "Team on port 4001 should be blocked"
|