litellm/tests/integration/management/test_team_delete_large_membership.py
yuneng-jiang 6530b329f4
fix(proxy): give /team/delete transactions a 30s timeout (#46979)
Both interactive transactions in /team/delete used Prisma's 5s default.
A lock wait or a short database stall past 5s expired the transaction
after the keys were already deleted, so the call returned 503 and left
the team row, memberships and team models behind
2026-10-16 14:36:48 -07:00

749 lines
33 KiB
Python

"""`/team/delete` as one locked transaction, whatever the roster size.
The delete removes the team row, its membership rows, every member's `teams` reference and every
team key in one pass, writes one tombstone per team, evicts the cached team object and takes the
team's advisory lock (the one `/team/member_add` takes) before it writes. A roster larger than the
Prisma pool used to fail with P2028 because each member got its own transaction.
"""
from __future__ import annotations
import json
import os
import uuid
from collections.abc import Callable, Mapping, Sequence
from concurrent.futures import Future, ThreadPoolExecutor
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import psycopg
import pytest
import yaml
from pydantic import JsonValue
from redis import Redis
from tests.integration._support.client import (
Gateway,
Scenario,
delete_key_if_present,
eventually,
object_value,
string_value,
)
from tests.integration._support.database import read_rows, write_rows
from tests.integration._support.process import owned_proxy_process
LARGE_ROSTER: Final = 250
POOL_LIMIT: Final = 5
# One statement seeds the whole roster: 250 individual /user/new calls would dominate the runtime.
SEED_USERS_SQL: Final = """
INSERT INTO "LiteLLM_UserTable" (user_id, user_role, teams, models)
SELECT %s || '-' || lpad(n::text, 3, '0'), 'internal_user', '{}'::text[], '{}'::text[]
FROM generate_series(1, %s::int) AS n
"""
# The master key's user id. `/team/new` appends the creator to the roster as an admin, so every team
# created here carries this member alongside the ones the test adds.
PROXY_ADMIN: Final = "default_user_id"
KEYS_WELL_PAST_THE_PUBLISH_CAP: Final = 2500
SEED_TEAM_KEYS_SQL: Final = """
INSERT INTO "LiteLLM_VerificationToken" (token, team_id, models)
SELECT encode(sha256(convert_to(%s || '-' || n::text, 'UTF8')), 'hex'), %s, '{}'::text[]
FROM generate_series(1, %s::int) AS n
"""
TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))"
LOCK_MODEL_ROW_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s FOR UPDATE'
LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS: Final = 7
OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """
SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age
FROM pg_stat_activity
WHERE %s::int = ANY(pg_blocking_pids(pid))
"""
# Sessions blocked on the advisory lock the given backend holds, and nothing else on the shared rig.
WAITERS_ON_HELD_LOCK_SQL: Final = """
SELECT count(*)::int AS waiting
FROM pg_locks waiter
JOIN pg_stat_activity session ON session.pid = waiter.pid
WHERE waiter.locktype = 'advisory'
AND NOT waiter.granted
AND session.wait_event_type = 'Lock'
AND session.query ILIKE %s
AND (waiter.classid, waiter.objid, waiter.objsubid) IN (
SELECT held.classid, held.objid, held.objsubid
FROM pg_locks held
WHERE held.locktype = 'advisory' AND held.granted AND held.pid = %s::int
)
"""
def _hashed(key: str) -> str:
return sha256(key.encode()).hexdigest()
def _team_rows(team_id: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
def _membership_user_ids(team_id: str) -> list[JsonValue]:
rows: Final = read_rows(
'SELECT user_id FROM "LiteLLM_TeamMembership" WHERE team_id = %s ORDER BY user_id', (team_id,)
)
return [row["user_id"] for row in rows]
def _user_teams(user_id: str) -> JsonValue:
rows: Final = read_rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (user_id,))
assert len(rows) == 1, f"user row for {user_id}: {rows}"
return rows[0]["teams"]
def _users_referencing(team_id: str) -> list[JsonValue]:
rows: Final = read_rows(
'SELECT user_id FROM "LiteLLM_UserTable" WHERE %s = ANY(teams) ORDER BY user_id', (team_id,)
)
return [row["user_id"] for row in rows]
def _user_ids_with_prefix(prefix: str) -> list[JsonValue]:
rows: Final = read_rows(
'SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id LIKE %s ORDER BY user_id', (f"{prefix}-%",)
)
return [row["user_id"] for row in rows]
def _live_token(hashed: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT token, team_id FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,))
def _deleted_token(hashed: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT token, team_id FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s', (hashed,))
def _tombstones(team_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT team_id, members_with_roles FROM "LiteLLM_DeletedTeamTable" WHERE team_id = %s', (team_id,)
)
def _roster_user_ids(roster: JsonValue) -> list[str]:
assert isinstance(roster, list), f"roster is not a list: {roster!r}"
return sorted(string_value(object_value(member)["user_id"]) for member in roster)
def _waiters_on_lock_held_by(backend_pid: int) -> int:
rows: Final = read_rows(WAITERS_ON_HELD_LOCK_SQL, ("%pg_advisory_xact_lock%", str(backend_pid)))
waiting: Final = rows[0]["waiting"]
assert isinstance(waiting, int)
return waiting
def _oldest_transaction_blocked_by(backend_pid: int) -> float:
rows: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(backend_pid),))
age: Final = rows[0]["age"]
assert isinstance(age, float)
return age
def _remove_team_by_sql(team_id: str) -> None:
"""Cleanup for a team the test expects to have deleted itself. Whatever a failed delete left behind
(row, memberships, `teams` references) goes by SQL so the shared rig stays clean without sending
another request through the proxy under test."""
if not _team_rows(team_id):
return
write_rows('DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = %s', (team_id,))
write_rows(
'UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, %s) WHERE %s = ANY(teams)', (team_id, team_id)
)
write_rows('DELETE FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
assert _team_rows(team_id) == []
def _remove_team_if_present(gateway: Gateway, team_id: str) -> None:
"""Cleanup for teams the test deletes itself: the API delete first, SQL for anything it leaves."""
if not _team_rows(team_id):
return
gateway.request("POST", "/team/delete", {"team_ids": [team_id]})
_remove_team_by_sql(team_id)
def _create_team(scenario: Scenario, **fields: JsonValue) -> str:
"""A team the test deletes itself; cleanup only removes it if the test left it behind."""
created: Final = scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", **fields})
team_id: Final = string_value(created["team_id"])
scenario.cleanups.callback(_remove_team_if_present, scenario.gateway, team_id)
return team_id
def _create_team_model(scenario: Scenario, team_id: str) -> str:
created: Final = scenario.gateway.post(
"/model/new",
{
"model_name": f"integration-{uuid.uuid4().hex}",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-integration"},
"model_info": {"team_id": team_id},
},
)
model_id: Final = string_value(object_value(created["model_info"])["id"])
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))
return model_id
def _team_model_rows(model_id: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))
def _delete_while_locked(gateway: Gateway, team_id: str, lock_sql: str, lock_target: str) -> httpx.Response:
with ThreadPoolExecutor(max_workers=1) as pool:
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(lock_sql, (lock_target,))
pending: Final = _post_in_background(pool, gateway, "/team/delete", {"team_ids": [team_id]})
eventually(
lambda: _oldest_transaction_blocked_by(holder.info.backend_pid),
lambda age: age > LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS,
seconds=30,
)
assert not pending.done(), "delete returned while the lock was still held"
return pending.result(timeout=60)
def _generate_key(scenario: Scenario, **fields: JsonValue) -> str:
"""A key the team delete is expected to remove; cleanup only deletes it if it is still live."""
key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"])
scenario.cleanups.callback(delete_key_if_present, scenario.gateway, key)
return key
def _delete_seeded_users(prefix: str) -> None:
write_rows('DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (f"{prefix}-%",))
assert _user_ids_with_prefix(prefix) == []
def _seed_users(scenario: Scenario, prefix: str, count: int) -> tuple[str, ...]:
"""Insert `count` user rows in one statement; ids are `<prefix>-001` … `<prefix>-<count>`."""
users: Final = tuple(f"{prefix}-{index:03d}" for index in range(1, count + 1))
write_rows(SEED_USERS_SQL, (prefix, str(count)))
scenario.cleanups.callback(_delete_seeded_users, prefix)
assert _user_ids_with_prefix(prefix) == list(users)
return users
def _bulk_member_add(gateway: Gateway, team_id: str, users: Sequence[str]) -> None:
gateway.post(
"/team/member_add",
{"team_id": team_id, "member": [{"role": "user", "user_id": user_id} for user_id in users]},
)
def _delete_teams(gateway: Gateway, team_ids: Sequence[str]) -> httpx.Response:
return gateway.request("POST", "/team/delete", {"team_ids": list(team_ids)})
def _team_info(gateway: Gateway, team_id: str) -> httpx.Response:
return gateway.request("GET", "/team/info", params={"team_id": team_id})
def _team_not_found_body(team_id: str) -> dict[str, JsonValue]:
"""The proxy's exception handler wraps the 404 detail as an `error` object with the detail stringified."""
return {
"error": {
"message": f"{{'message': 'Team not found, passed team id: {team_id}.'}}",
"type": "auth_error",
"param": "None",
"code": "404",
}
}
def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"team delete {uuid.uuid4().hex}"}]},
key=key,
)
def _post_with_timeout(gateway: Gateway, path: str, body: Mapping[str, JsonValue], timeout: float) -> httpx.Response:
"""Like `Gateway.request` with a per-call timeout longer than the client's default 15 s."""
return gateway.client.request(
"POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}, timeout=timeout
)
def _post_in_background(
pool: ThreadPoolExecutor, gateway: Gateway, path: str, body: Mapping[str, JsonValue]
) -> Future[httpx.Response]:
return pool.submit(_post_with_timeout, gateway, path, body, 60)
def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"]["database_connection_pool_limit"] = pool_limit
config["general_settings"]["database_connection_pool_timeout"] = 60
path: Final = tmp_path / f"pool-{pool_limit}.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _redis() -> Redis:
return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
def test_delete_small_team_removes_rows_keys_tombstone_and_cache(
gateway: Gateway, record_property: Callable[[str, object], None]
) -> None:
with gateway.scenario() as scenario, _redis() as cache:
model: Final = scenario.model()
users: Final = sorted(scenario.user() for _ in range(3))
team: Final = _create_team(scenario)
_bulk_member_add(gateway, team, users)
keys: Final = tuple(_generate_key(scenario, team_id=team) for _ in range(2))
hashed: Final = tuple(_hashed(key) for key in keys)
roster: Final = sorted([PROXY_ADMIN, *users])
assert _membership_user_ids(team) == roster
assert _users_referencing(team) == roster
assert all(_user_teams(user) == [team] for user in users), [_user_teams(user) for user in users]
assert all(len(_live_token(digest)) == 1 for digest in hashed), hashed
warm: Final = _chat(gateway, model, keys[0])
assert warm.status_code == 200, warm.text
team_cache_key: Final = f"team_id:{team}"
eventually(lambda: cache.exists(team_cache_key), lambda present: present == 1, seconds=10)
record_property("redis_keys_before_delete", sorted(entry.decode() for entry in cache.keys(f"*{team}*")))
response: Final = _delete_teams(gateway, [team])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _users_referencing(team) == []
assert [_user_teams(user) for user in users] == [[], [], []]
assert [_live_token(digest) for digest in hashed] == [[], []]
assert [_deleted_token(digest) for digest in hashed] == [
[{"token": hashed[0], "team_id": team}],
[{"token": hashed[1], "team_id": team}],
]
tombstones: Final = _tombstones(team)
assert len(tombstones) == 1, tombstones
assert tombstones[0]["team_id"] == team
assert _roster_user_ids(tombstones[0]["members_with_roles"]) == roster
info: Final = _team_info(gateway, team)
assert info.status_code == 404, info.text
assert info.json() == _team_not_found_body(team)
assert cache.exists(team_cache_key) == 0
record_property("redis_keys_after_delete", sorted(entry.decode() for entry in cache.keys(f"*{team}*")))
@pytest.mark.timeout(240) # owned two-worker proxy boot plus a 250-member roster
def test_delete_250_member_team_succeeds_with_pool_limit_five_on_two_workers(
gateway: Gateway, tmp_path: Path, record_property: Callable[[str, object], None]
) -> None:
prefix: Final = f"integration-roster-{uuid.uuid4().hex}"
# The scenario is bound to the shared gateway and its cleanups are SQL, so a failed delete on the
# owned proxy (and whatever it does to that proxy's workers) cannot mask the assertion below with a
# second failure during cleanup. The owned proxy is stopped before the cleanups run.
with (
gateway.scenario() as scenario,
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"]},
config=_config_with_pool_limit(tmp_path, POOL_LIMIT),
remove_environment=("DATABASE_URL_READ_REPLICA",),
workers=2,
) as owned,
):
team: Final = string_value(
owned.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"]
)
scenario.cleanups.callback(_remove_team_by_sql, team)
users: Final = _seed_users(scenario, prefix, LARGE_ROSTER)
added: Final = _post_with_timeout(
owned.gateway,
"/team/member_add",
{"team_id": team, "member": [{"role": "user", "user_id": user_id} for user_id in users]},
timeout=120,
)
assert added.status_code == 200, added.text
roster: Final = sorted([PROXY_ADMIN, *users])
assert _membership_user_ids(team) == roster
assert _users_referencing(team) == roster
response: Final = _post_with_timeout(owned.gateway, "/team/delete", {"team_ids": [team]}, timeout=120)
record_property("h2_delete_response", f"{response.status_code} {response.text[:300]}")
assert response.status_code == 200, (
f"/team/delete of a {LARGE_ROSTER}-member team with database_connection_pool_limit={POOL_LIMIT}: "
f"{response.status_code} {response.text}"
)
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _users_referencing(team) == []
assert _user_ids_with_prefix(prefix) == list(users)
tombstones: Final = _tombstones(team)
assert len(tombstones) == 1, tombstones
assert _roster_user_ids(tombstones[0]["members_with_roles"]) == roster
def test_delete_waits_for_the_team_advisory_lock_and_completes_after_release(
gateway: Gateway, record_property: Callable[[str, object], None]
) -> None:
with gateway.scenario() as scenario:
user: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(gateway, team, [user])
hashed: Final = _hashed(_generate_key(scenario, team_id=team))
with ThreadPoolExecutor(max_workers=1) as pool:
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(TAKE_TEAM_LOCK_SQL, (team,))
pending: Final = _post_in_background(pool, gateway, "/team/delete", {"team_ids": [team]})
eventually(
lambda: _waiters_on_lock_held_by(holder.info.backend_pid),
lambda waiting: waiting >= 1,
seconds=20,
)
assert not pending.done(), "delete returned while the team lock was still held"
assert _team_rows(team) == [{"team_id": team}], "team row deleted while the team lock was held"
# Recorded before the count assertion so both legs document what the delete had already
# written by the time it reached the lock.
record_property(
"state_while_blocked",
json.dumps(
{
"membership_user_ids": _membership_user_ids(team),
"user_teams": _user_teams(user),
"live_token_rows": len(_live_token(hashed)),
"tombstones": len(_tombstones(team)),
}
),
)
# One transaction per delete: a per-member fan-out would queue one waiter per roster entry.
eventually(
lambda: _waiters_on_lock_held_by(holder.info.backend_pid),
lambda waiting: waiting == 1,
seconds=10,
)
# leaving the holder block commits its transaction, which releases the advisory lock
response: Final = pending.result(timeout=60)
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _user_teams(user) == []
assert _live_token(hashed) == []
assert len(_tombstones(team)) == 1, _tombstones(team)
def test_delete_completes_after_waiting_on_a_locked_team_model_row_past_the_default_transaction_timeout(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
team: Final = _create_team(scenario)
model_id: Final = _create_team_model(scenario, team)
response: Final = _delete_while_locked(gateway, team, LOCK_MODEL_ROW_SQL, model_id)
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _team_model_rows(model_id) == []
assert len(_tombstones(team)) == 1, _tombstones(team)
def test_delete_completes_after_waiting_on_the_team_lock_past_the_default_transaction_timeout(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
user: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(gateway, team, [user])
response: Final = _delete_while_locked(gateway, team, TAKE_TEAM_LOCK_SQL, team)
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _user_teams(user) == []
assert len(_tombstones(team)) == 1, _tombstones(team)
def test_deleting_two_teams_sharing_a_member_in_one_call_clears_both_from_the_member(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
user: Final = scenario.user()
first: Final = _create_team(scenario)
second: Final = _create_team(scenario)
_bulk_member_add(gateway, first, [user])
_bulk_member_add(gateway, second, [user])
assert _user_teams(user) == [first, second]
response: Final = _delete_teams(gateway, [first, second])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [first, second]}
assert _team_rows(first) == []
assert _team_rows(second) == []
assert _membership_user_ids(first) == []
assert _membership_user_ids(second) == []
assert _user_teams(user) == []
assert [row["team_id"] for row in _tombstones(first)] == [first]
assert [row["team_id"] for row in _tombstones(second)] == [second]
def test_deleting_one_team_leaves_the_members_other_team_and_key_intact(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
user: Final = scenario.user()
deleted: Final = _create_team(scenario)
kept: Final = scenario.team()
_bulk_member_add(gateway, deleted, [user])
_bulk_member_add(gateway, kept, [user])
kept_key: Final = scenario.key(team_id=kept, user_id=user)
before: Final = _chat(gateway, model, kept_key)
assert before.status_code == 200, before.text
response: Final = _delete_teams(gateway, [deleted])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [deleted]}
assert _team_rows(deleted) == []
assert _team_rows(kept) == [{"team_id": kept}]
assert _membership_user_ids(deleted) == []
assert _membership_user_ids(kept) == [PROXY_ADMIN, user]
assert _user_teams(user) == [kept]
assert _live_token(_hashed(kept_key)) == [{"token": _hashed(kept_key), "team_id": kept}]
after: Final = _chat(gateway, model, kept_key)
assert after.status_code == 200, after.text
@pytest.mark.timeout(240) # owned proxy boot
def test_delete_writes_one_audit_row_for_the_team_and_one_per_key(gateway: Gateway, tmp_path: Path) -> None:
with (
owned_proxy_process(
gateway,
tmp_path,
{"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_STORE_AUDIT_LOGS": "true"},
remove_environment=("DATABASE_URL_READ_REPLICA",),
) as owned,
owned.gateway.scenario() as scenario,
):
user: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(owned.gateway, team, [user])
hashed: Final = _hashed(_generate_key(scenario, team_id=team))
response: Final = _delete_teams(owned.gateway, [team])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
def deleted_audit_rows() -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT table_name, action, object_id FROM "LiteLLM_AuditLog" '
"WHERE object_id IN (%s, %s) AND action = 'deleted' ORDER BY table_name",
(team, hashed),
)
rows: Final = eventually(deleted_audit_rows, lambda found: len(found) >= 2, seconds=30)
assert rows == [
{"table_name": "LiteLLM_TeamTable", "action": "deleted", "object_id": team},
{"table_name": "LiteLLM_VerificationToken", "action": "deleted", "object_id": hashed},
]
def test_second_delete_of_the_same_team_is_404_with_one_tombstone(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
user: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(gateway, team, [user])
first: Final = _delete_teams(gateway, [team])
assert first.status_code == 200, first.text
assert first.json() == {"deleted_teams": [team]}
second: Final = _delete_teams(gateway, [team])
assert second.status_code == 404, second.text
assert second.json() == {"detail": {"error": f"Team not found, passed team_id={team}"}}
assert _team_rows(team) == []
assert [row["team_id"] for row in _tombstones(team)] == [team]
assert _user_teams(user) == []
def test_delete_empty_team_writes_tombstone_and_team_info_is_404(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
team: Final = _create_team(scenario)
present: Final = _team_info(gateway, team)
assert present.status_code == 200, present.text
response: Final = _delete_teams(gateway, [team])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _tombstones(team) == [
{"team_id": team, "members_with_roles": [{"role": "admin", "user_id": PROXY_ADMIN, "user_email": None}]}
]
info: Final = _team_info(gateway, team)
assert info.status_code == 404, info.text
assert info.json() == _team_not_found_body(team)
def test_delete_keys_only_team_removes_keys_and_revokes_them(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
model: Final = scenario.model()
team: Final = _create_team(scenario)
keys: Final = tuple(_generate_key(scenario, team_id=team) for _ in range(2))
hashed: Final = tuple(_hashed(key) for key in keys)
assert _membership_user_ids(team) == [PROXY_ADMIN]
for key in keys:
warm = _chat(gateway, model, key)
assert warm.status_code == 200, warm.text
response: Final = _delete_teams(gateway, [team])
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": [team]}
assert _team_rows(team) == []
assert [_live_token(digest) for digest in hashed] == [[], []]
assert [_deleted_token(digest) for digest in hashed] == [
[{"token": hashed[0], "team_id": team}],
[{"token": hashed[1], "team_id": team}],
]
for key in keys:
revoked = _chat(gateway, model, key)
assert revoked.status_code == 401, f"{revoked.status_code} {revoked.text}"
def test_delete_three_teams_in_one_call_lists_all_and_tombstones_each_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
teams: Final = tuple(_create_team(scenario) for _ in range(3))
for team in teams:
_bulk_member_add(gateway, team, [scenario.user()])
response: Final = _delete_teams(gateway, teams)
assert response.status_code == 200, response.text
assert response.json() == {"deleted_teams": list(teams)}
for team in teams:
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _users_referencing(team) == []
assert [row["team_id"] for row in _tombstones(team)] == [team]
def test_recreating_the_same_team_id_after_delete_serves_the_fresh_team(gateway: Gateway) -> None:
with gateway.scenario() as scenario, _redis() as cache:
model: Final = scenario.model()
original_member: Final = scenario.user()
replacement_member: Final = scenario.user()
team: Final = _create_team(scenario)
_bulk_member_add(gateway, team, [original_member])
original_key: Final = _generate_key(scenario, team_id=team)
warm: Final = _chat(gateway, model, original_key)
assert warm.status_code == 200, warm.text
team_cache_key: Final = f"team_id:{team}"
eventually(lambda: cache.exists(team_cache_key), lambda present: present == 1, seconds=10)
response: Final = _delete_teams(gateway, [team])
assert response.status_code == 200, response.text
assert _team_rows(team) == []
assert cache.exists(team_cache_key) == 0
fresh_alias: Final = f"integration-recreated-{uuid.uuid4().hex}"
recreated: Final = gateway.request(
"POST",
"/team/new",
{
"team_id": team,
"team_alias": fresh_alias,
"members_with_roles": [{"role": "user", "user_id": replacement_member}],
},
)
assert recreated.status_code == 200, recreated.text
assert recreated.json()["team_id"] == team
info: Final = _team_info(gateway, team)
assert info.status_code == 200, info.text
team_info: Final = object_value(info.json()["team_info"])
assert team_info["team_alias"] == fresh_alias
assert _roster_user_ids(team_info["members_with_roles"]) == [PROXY_ADMIN, replacement_member]
assert _membership_user_ids(team) == [PROXY_ADMIN, replacement_member]
assert _user_teams(replacement_member) == [team]
assert _user_teams(original_member) == []
fresh_key: Final = scenario.key(team_id=team)
served: Final = _chat(gateway, model, fresh_key)
assert served.status_code == 200, served.text
cached: Final = eventually(lambda: cache.get(team_cache_key), lambda value: value is not None, seconds=10)
assert isinstance(cached, bytes), cached
assert json.loads(cached)["team_alias"] == fresh_alias, cached
revoked: Final = _chat(gateway, model, original_key)
assert revoked.status_code == 401, f"{revoked.status_code} {revoked.text}"
def test_member_add_and_delete_released_together_leave_no_team_reference(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
newcomer: Final = scenario.user()
team: Final = _create_team(scenario)
with ThreadPoolExecutor(max_workers=2) as pool:
with psycopg.connect(os.environ["DATABASE_URL"]) as holder:
holder.execute(TAKE_TEAM_LOCK_SQL, (team,))
pending_delete: Final = _post_in_background(pool, gateway, "/team/delete", {"team_ids": [team]})
pending_add: Final = _post_in_background(
pool,
gateway,
"/team/member_add",
{"team_id": team, "member": {"role": "user", "user_id": newcomer}},
)
eventually(
lambda: _waiters_on_lock_held_by(holder.info.backend_pid),
lambda waiting: waiting == 2,
seconds=20,
)
assert not pending_delete.done() and not pending_add.done()
# leaving the holder block commits its transaction, which releases the advisory lock
deleted: Final = pending_delete.result(timeout=60)
added: Final = pending_add.result(timeout=60)
assert deleted.status_code == 200, deleted.text
assert deleted.json() == {"deleted_teams": [team]}
assert added.status_code in (200, 404), f"{added.status_code} {added.text}"
assert _team_rows(team) == []
assert _membership_user_ids(team) == []
assert _user_teams(newcomer) == []
assert _users_referencing(team) == []
assert [row["team_id"] for row in _tombstones(team)] == [team]
def _seeded_team_keys(team_id: str, count: int) -> tuple[str, ...]:
prefix: Final = f"sk-integration-{uuid.uuid4().hex}"
write_rows(SEED_TEAM_KEYS_SQL, (prefix, team_id, str(count)))
return tuple(f"{prefix}-{index}" for index in range(1, count + 1))
def _models_status(proxy: Gateway, key: str) -> int:
return proxy.request("GET", "/v1/models", key=key).status_code
def _statuses(proxy: Gateway, keys: Sequence[str]) -> dict[int, int]:
with ThreadPoolExecutor(max_workers=16) as pool:
statuses: Final = tuple(pool.map(lambda key: _models_status(proxy, key), keys))
return {status: statuses.count(status) for status in frozenset(statuses)}
def test_deleting_a_team_with_more_keys_than_the_publish_cap_revokes_every_key_on_a_warmed_peer(
gateway: Gateway, peer: Gateway
) -> None:
with gateway.scenario() as scenario:
team: Final = _create_team(scenario)
scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,))
keys: Final = _seeded_team_keys(team, KEYS_WELL_PAST_THE_PUBLISH_CAP)
assert _statuses(peer, keys) == {200: len(keys)}
response: Final = _delete_teams(gateway, [team])
assert response.status_code == 200, response.text
assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) == []
_ = eventually(lambda: _models_status(peer, keys[-1]), lambda status: status == 401, seconds=10)
assert _statuses(peer, keys) == {401: len(keys)}