mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test(integration): cache the team member default budget in Redis as JSON (Pylon #7180)
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
31601828c1
commit
8efad109b4
2 changed files with 43 additions and 0 deletions
|
|
@ -1779,6 +1779,9 @@
|
|||
],
|
||||
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
|
||||
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
|
||||
],
|
||||
"tests/integration/management/test_team_member_budget_cache.py::test_team_member_default_budget_lands_in_redis_after_first_member_call": [
|
||||
"mgmt.team_member_budget.default_budget_is_cached_in_redis_as_json"
|
||||
]
|
||||
},
|
||||
"browser": {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,40 @@
|
|||
import os
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from redis import Redis
|
||||
|
||||
from tests.integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from tests.integration._support.database import read_rows
|
||||
|
||||
_CACHED_BUDGET: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@pytest.mark.covers("mgmt.team_member_budget.default_budget_is_cached_in_redis_as_json")
|
||||
def test_team_member_default_budget_lands_in_redis_after_first_member_call(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
user: Final = scenario.user()
|
||||
team: Final = scenario.team(team_member_budget=25)
|
||||
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
|
||||
teams: Final = read_rows('SELECT metadata FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
|
||||
assert len(teams) == 1, teams
|
||||
budget_id: Final = string_value(object_value(teams[0]["metadata"])["team_member_budget_id"])
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "member budget cache"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
||||
cached: Final = eventually(
|
||||
lambda: cache.get(f"team_member_default_budget:{budget_id}"),
|
||||
lambda value: value is not None,
|
||||
seconds=10,
|
||||
)
|
||||
assert isinstance(cached, bytes), cached
|
||||
budget: Final = _CACHED_BUDGET.validate_json(cached)
|
||||
assert budget["budget_id"] == budget_id, cached
|
||||
assert budget["max_budget"] == 25, cached
|
||||
Loading…
Add table
Reference in a new issue