mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
93b29c3977
* fix(proxy): broadcast team cache invalidation on team writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover cross-instance team cache invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
from typing import Final
|
|
|
|
import httpx
|
|
|
|
from tests.integration._support.client import Gateway, eventually, object_value
|
|
|
|
_PROMPT: Final = {"messages": [{"role": "user", "content": "team blocking"}]}
|
|
|
|
|
|
def _chat(proxy: Gateway, model: str, key: str) -> httpx.Response:
|
|
return proxy.request("POST", "/v1/chat/completions", {"model": model, **_PROMPT}, key=key)
|
|
|
|
|
|
def _rejected(proxy: Gateway, model: str, key: str) -> httpx.Response:
|
|
return eventually(lambda: _chat(proxy, model, key), lambda response: response.status_code == 401, seconds=10)
|
|
|
|
|
|
def _blocked_flag(proxy: Gateway, team_id: str) -> object:
|
|
return object_value(proxy.get("/team/info", {"team_id": team_id})["team_info"])["blocked"]
|
|
|
|
|
|
def test_a_team_blocked_on_one_instance_is_rejected_on_both(gateway: Gateway, peer: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
team: Final = scenario.team()
|
|
key: Final = scenario.key(team_id=team)
|
|
assert _blocked_flag(peer, team) is False
|
|
allowed: Final = _chat(gateway, model, key)
|
|
assert allowed.status_code == 200, allowed.text
|
|
peer_allowed: Final = _chat(peer, model, key)
|
|
assert peer_allowed.status_code == 200, peer_allowed.text
|
|
|
|
peer.post("/team/update", {"team_id": team, "blocked": True})
|
|
|
|
assert _blocked_flag(gateway, team) is True
|
|
assert _blocked_flag(peer, team) is True
|
|
rejections: Final = tuple(_rejected(proxy, model, key) for proxy in (gateway, peer))
|
|
|
|
for rejection in rejections:
|
|
error: Final = object_value(rejection.json()["error"])
|
|
assert error["type"] == "auth_error", rejection.text
|
|
assert f"Team={team} is blocked" in str(error["message"]), rejection.text
|