mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
340 lines
15 KiB
Python
340 lines
15 KiB
Python
import json
|
|
import os
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import Final
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
|
|
import httpx
|
|
import psycopg
|
|
import pytest
|
|
from integration._support.client import Gateway, delete_key_if_present, eventually, string_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
from psycopg import sql
|
|
from pydantic import BaseModel, JsonValue, TypeAdapter
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
|
|
@pytest.mark.covers("other.database.regeneration.writer_updates_dependent_grants")
|
|
def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gateway, tmp_path: Path) -> None:
|
|
role: Final = f"integration_reader_{uuid.uuid4().hex}"
|
|
url: Final = os.environ["DATABASE_URL"]
|
|
parsed: Final = urlsplit(url)
|
|
reader_url: Final = urlunsplit(
|
|
parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")
|
|
)
|
|
with psycopg.connect(url, autocommit=True) as admin:
|
|
admin.execute(
|
|
sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(
|
|
sql.Identifier(role)
|
|
)
|
|
)
|
|
try:
|
|
admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role)))
|
|
admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role)))
|
|
admin.execute(sql.SQL("ALTER ROLE {} SET default_transaction_read_only = on").format(sql.Identifier(role)))
|
|
with psycopg.connect(reader_url, autocommit=True) as reader:
|
|
assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",)
|
|
with pytest.raises(psycopg.errors.ReadOnlySqlTransaction):
|
|
reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false')
|
|
with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate:
|
|
assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), (
|
|
"Candidate reader was never connected"
|
|
)
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
outside: Final = scenario.model()
|
|
old: Final = string_value(candidate.post("/key/generate", {"models": [outside]})["key"])
|
|
new: Final = f"sk-integration-{uuid.uuid4().hex}"
|
|
scenario.cleanups.callback(delete_key_if_present, gateway, old)
|
|
scenario.cleanups.callback(delete_key_if_present, gateway, new)
|
|
old_hash: Final = sha256(old.encode()).hexdigest()
|
|
before: Final = candidate.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": "no grant yet"}]},
|
|
key=old,
|
|
)
|
|
assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", (
|
|
before.text
|
|
)
|
|
response: Final = candidate.request(
|
|
"POST",
|
|
"/v1/access_group",
|
|
{
|
|
"access_group_name": f"integration-{uuid.uuid4().hex}",
|
|
"access_model_names": [model],
|
|
"assigned_key_ids": [old_hash],
|
|
},
|
|
)
|
|
assert response.status_code == 201, response.text
|
|
group: Final = string_value(response.json()["access_group_id"])
|
|
try:
|
|
with psycopg.connect(url) as blocker, ThreadPoolExecutor(max_workers=1) as executor:
|
|
blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE')
|
|
pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}")
|
|
try:
|
|
reached: Final = eventually(
|
|
lambda: read_rows(
|
|
"SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) "
|
|
"AND usename=%s AND query LIKE 'SELECT%%'",
|
|
(blocker.info.backend_pid, role),
|
|
),
|
|
bool,
|
|
seconds=3,
|
|
)
|
|
assert reached == [{"usename": role}]
|
|
finally:
|
|
blocker.rollback()
|
|
selected: Final = pending.result(timeout=5)
|
|
assert selected.status_code == 200 and selected.json()["access_group_id"] == group, (
|
|
selected.text
|
|
)
|
|
assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40
|
|
regenerated: Final = candidate.post(
|
|
"/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}
|
|
)
|
|
assert regenerated["key"] == new
|
|
new_hash: Final = sha256(new.encode()).hexdigest()
|
|
assert new != old
|
|
assert read_rows(
|
|
'SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)
|
|
) == [{"assigned_key_ids": [new_hash]}]
|
|
assert read_rows(
|
|
'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)',
|
|
([old_hash, new_hash],),
|
|
) == [{"token": new_hash, "access_group_ids": [group]}]
|
|
assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40
|
|
assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40
|
|
denied: Final = candidate.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": "rotated key"}]},
|
|
key=old,
|
|
)
|
|
assert (
|
|
denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db"
|
|
), denied.text
|
|
finally:
|
|
deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}")
|
|
assert deleted.status_code == 204, deleted.text
|
|
assert (
|
|
read_rows(
|
|
'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s',
|
|
(group,),
|
|
)
|
|
== []
|
|
)
|
|
finally:
|
|
admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
|
|
admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role)))
|
|
assert read_rows("SELECT rolname FROM pg_roles WHERE rolname=%s", (role,)) == []
|
|
|
|
|
|
class _Usage(BaseModel):
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
|
|
|
|
class _AssistantMessage(BaseModel):
|
|
content: str
|
|
|
|
|
|
class _Choice(BaseModel):
|
|
message: _AssistantMessage
|
|
|
|
|
|
class _ChatResponse(BaseModel):
|
|
id: str
|
|
choices: tuple[_Choice, ...]
|
|
usage: _Usage
|
|
|
|
|
|
class _GenerateKeyResponse(BaseModel):
|
|
key: str
|
|
|
|
|
|
class _Error(BaseModel):
|
|
type: str
|
|
|
|
|
|
class _ErrorResponse(BaseModel):
|
|
error: _Error
|
|
|
|
|
|
class _SpendLog(BaseModel):
|
|
api_key: str
|
|
spend: float
|
|
|
|
|
|
class _KeySpend(BaseModel):
|
|
spend: float
|
|
|
|
|
|
class _DeprecatedKey(BaseModel):
|
|
active_token_id: str
|
|
grace_seconds: float
|
|
|
|
|
|
_CHAT_BODY: Final = TypeAdapter(dict[str, JsonValue])
|
|
_EXPECTED_CHAT: Final = ("POST", "/v1/chat/completions")
|
|
|
|
|
|
def _chat_body(model: str, prompt: str) -> dict[str, JsonValue]:
|
|
return {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
}
|
|
|
|
|
|
def _wire_reply(identity: str) -> Reply:
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"id": identity,
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": "integration-keys-auth-grace",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": "wire response"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
|
|
def _chat(gateway: Gateway, model: str, key: str, prompt: str) -> httpx.Response:
|
|
return gateway.request("POST", "/v1/chat/completions", _chat_body(model, prompt), key=key)
|
|
|
|
|
|
def _assert_success(response: httpx.Response) -> _ChatResponse:
|
|
assert response.status_code == 200, response.text
|
|
payload: Final = _ChatResponse.model_validate_json(response.content)
|
|
assert payload.choices[0].message.content == "wire response", response.text
|
|
assert payload.usage == _Usage(prompt_tokens=20, completion_tokens=20, total_tokens=40), response.text
|
|
return payload
|
|
|
|
|
|
@pytest.mark.parametrize("spelling", ("body", "path", "plaintext_path"))
|
|
def test_regenerate_with_a_grace_period_keeps_the_old_key_serving_billed_to_the_new_row(
|
|
gateway: Gateway, spelling: str
|
|
) -> None:
|
|
assert os.environ.get("LITELLM_KEY_ROTATION_GRACE_PERIOD") is None
|
|
backend: Final = "integration-keys-auth-grace"
|
|
prompts: Final = (
|
|
"grace old key " + uuid.uuid4().hex,
|
|
"grace new key " + uuid.uuid4().hex,
|
|
"revoked second key " + uuid.uuid4().hex,
|
|
"final generated key " + uuid.uuid4().hex,
|
|
)
|
|
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/v1/chat/completions"
|
|
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
|
|
body: Final = _CHAT_BODY.validate_json(request.body)
|
|
assert body in tuple(_chat_body(backend, prompt) for prompt in prompts), body
|
|
return _wire_reply("chatcmpl-" + uuid.uuid4().hex)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(
|
|
model=f"openai/{backend}",
|
|
api_base=f"{wire.url}/v1",
|
|
api_key="synthetic-openai-key",
|
|
input_cost_per_token=0.001,
|
|
output_cost_per_token=0.002,
|
|
)
|
|
created: Final = gateway.request("POST", "/key/generate", {"models": [model]})
|
|
assert created.status_code == 200, created.text
|
|
key_one: Final = _GenerateKeyResponse.model_validate_json(created.content).key
|
|
scenario.cleanups.callback(delete_key_if_present, gateway, key_one)
|
|
hash_one: Final = sha256(key_one.encode()).hexdigest()
|
|
|
|
path_key_one: Final = hash_one if spelling == "path" else key_one
|
|
generated_two: Final = (
|
|
gateway.request("POST", "/key/regenerate", {"key": key_one, "grace_period": "1h"})
|
|
if spelling == "body"
|
|
else gateway.request("POST", f"/key/{path_key_one}/regenerate", {"grace_period": "1h"})
|
|
)
|
|
assert generated_two.status_code == 200, generated_two.text
|
|
key_two: Final = _GenerateKeyResponse.model_validate_json(generated_two.content).key
|
|
scenario.cleanups.callback(delete_key_if_present, gateway, key_two)
|
|
hash_two: Final = sha256(key_two.encode()).hexdigest()
|
|
old_response: Final = _assert_success(_chat(gateway, model, key_one, prompts[0]))
|
|
new_response: Final = _assert_success(_chat(gateway, model, key_two, prompts[1]))
|
|
assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,) * 2
|
|
|
|
old_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
|
|
eventually(
|
|
lambda: read_rows(
|
|
'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(old_response.id,),
|
|
),
|
|
lambda rows: len(rows) == 1,
|
|
seconds=70,
|
|
)
|
|
)
|
|
new_spend: Final = TypeAdapter(tuple[_SpendLog, ...]).validate_python(
|
|
eventually(
|
|
lambda: read_rows(
|
|
'SELECT api_key, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(new_response.id,),
|
|
),
|
|
lambda rows: len(rows) == 1,
|
|
seconds=70,
|
|
)
|
|
)
|
|
assert old_spend[0].api_key == hash_two
|
|
assert old_spend[0].spend > 0
|
|
assert new_spend[0].api_key == hash_two
|
|
assert new_spend[0].spend > 0
|
|
key_spend: Final = TypeAdapter(tuple[_KeySpend, ...]).validate_python(
|
|
eventually(
|
|
lambda: read_rows(
|
|
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s',
|
|
(hash_two,),
|
|
),
|
|
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
|
|
seconds=70,
|
|
)
|
|
)
|
|
assert key_spend[0].spend > 0
|
|
deprecated: Final = TypeAdapter(tuple[_DeprecatedKey, ...]).validate_python(
|
|
read_rows(
|
|
"SELECT active_token_id, EXTRACT(EPOCH FROM revoke_at - now())::float8 AS grace_seconds "
|
|
'FROM "LiteLLM_DeprecatedVerificationToken" WHERE token=%s',
|
|
(hash_one,),
|
|
)
|
|
)
|
|
assert len(deprecated) == 1
|
|
assert deprecated[0].active_token_id == hash_two
|
|
assert deprecated[0].grace_seconds == pytest.approx(3600, abs=120)
|
|
|
|
path_key_two: Final = hash_two if spelling == "path" else key_two
|
|
generated_three: Final = (
|
|
gateway.request("POST", "/key/regenerate", {"key": key_two})
|
|
if spelling == "body"
|
|
else gateway.request("POST", f"/key/{path_key_two}/regenerate", {})
|
|
)
|
|
assert generated_three.status_code == 200, generated_three.text
|
|
key_three: Final = _GenerateKeyResponse.model_validate_json(generated_three.content).key
|
|
scenario.cleanups.callback(delete_key_if_present, gateway, key_three)
|
|
refused: Final = _chat(gateway, model, key_two, prompts[2])
|
|
assert refused.status_code == 401, refused.text
|
|
assert _ErrorResponse.model_validate_json(refused.content).error.type == "token_not_found_in_db", refused.text
|
|
assert wire.drain() == (), "the revoked second key reached the upstream"
|
|
assert (
|
|
_assert_success(_chat(gateway, model, key_three, prompts[3])).choices[0].message.content == "wire response"
|
|
)
|
|
assert tuple((request.method, request.target) for request in wire.drain()) == (_EXPECTED_CHAT,)
|