diff --git a/.github/workflows/test-e2e-fail-closed-team-member-budget.yml b/.github/workflows/test-e2e-fail-closed-team-member-budget.yml new file mode 100644 index 00000000000..f4446bb313b --- /dev/null +++ b/.github/workflows/test-e2e-fail-closed-team-member-budget.yml @@ -0,0 +1,216 @@ +name: "Fail-closed Team Member Budget E2E" + +on: + pull_request: + paths: + - litellm/proxy/auth/auth_checks.py + - tests/e2e/conftest.py + - tests/e2e/pytest.ini + - tests/e2e/gateway/fail_closed_team_member_budget_ci_config.yml + - tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py + - .github/workflows/test-e2e-fail-closed-team-member-budget.yml + workflow_dispatch: + +concurrency: + group: e2e-fail-closed-team-member-budget-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: {} + +jobs: + run: + runs-on: ubuntu-latest + timeout-minutes: 45 + environment: e2e-changed + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository + permissions: + contents: read + id-token: write + services: + postgres: + image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + valkey: + image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa + ports: + - 6379:6379 + options: >- + --health-cmd "valkey-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-fail-closed-budget-e2e + LITELLM_LOG: WARNING + JSON_LOGS: "true" + steps: + - name: Validate configuration + env: + ROLE: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} + run: test -n "${ROLE}" || { echo "::error::Set repo variable E2E_AWS_ROLE_TO_ASSUME to an OIDC role with read access to the e2e secrets"; exit 1; } + + - name: Checkout + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + ref: ${{ github.sha }} + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.13" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen \ + --extra proxy --extra proxy-runtime --extra extra_proxy \ + --extra semantic-router --extra bedrock-realtime \ + --group ci --group proxy-dev --group e2e-dev + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Configure AWS credentials + id: aws + uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0 + with: + role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} + aws-region: us-east-1 + role-session-name: litellm-e2e-fail-closed-budget-${{ github.run_id }} + role-duration-seconds: 900 + output-env-credentials: false + output-credentials: true + + - name: Fetch provider credentials from AWS Secrets Manager + env: + AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }} + AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }} + AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }} + AWS_DEFAULT_REGION: us-east-1 + run: | + umask 077 + aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-provider-keys \ + --query SecretString --output text \ + | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env + aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \ + --query SecretString --output text \ + | jq -R -s '{"LITELLM_LICENSE": .}' \ + | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env + printf 'DATABASE_URL=%s\nLITELLM_MASTER_KEY=%s\n' "${DATABASE_URL}" "${LITELLM_MASTER_KEY}" \ + > "${RUNNER_TEMP}/fail-closed-budget-values.env" + + - name: Start the fail-closed proxy + id: proxy + run: | + umask 077 + proxy_database_url="${DATABASE_URL}" + proxy_master_key="${LITELLM_MASTER_KEY}" + set -a + source tests/e2e/.env + set +a + export DATABASE_URL="${proxy_database_url}" + export LITELLM_MASTER_KEY="${proxy_master_key}" + nohup uv run --no-sync litellm \ + --config tests/e2e/gateway/fail_closed_team_member_budget_ci_config.yml --port 4000 \ + > "${RUNNER_TEMP}/fail-closed-budget-proxy.log" 2>&1 & + echo "E2E_PROXY_PID=$!" >> "${GITHUB_ENV}" + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "::error::fail-closed proxy did not become live" + exit 1 + + - name: Run the missing-membership e2e regression + id: e2e + env: + E2E_FAIL_CLOSED_BUDGET_STACK: "1" + E2E_FIXTURE_MODE: live + LITELLM_PROXY_URL: http://localhost:4000 + REDIS_HOST: 127.0.0.1 + REDIS_PORT: "6379" + run: | + umask 077 + report="${RUNNER_TEMP}/fail-closed-budget-e2e.xml" + log="${RUNNER_TEMP}/fail-closed-budget-e2e.log" + set +e + uv run --no-sync pytest tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py \ + -k test_missing_membership_counts_as_verified_zero_spend -v --reruns 0 -p no:cacheprovider \ + -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1 + status=$? + if [ -f "${report}" ]; then + uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \ + "${report}" tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py + verified=$? + else + verified=1 + fi + set -e + grep -E '^(FAILED|ERROR) ' "${log}" || true + grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1 + if [ "${status}" != "0" ] || [ "${verified}" != "0" ]; then + echo "::error::fail-closed team-member budget e2e regression failed or did not run" + exit 1 + fi + + - name: Redact e2e and proxy output + if: always() && steps.proxy.outcome != 'skipped' + run: | + umask 077 + files=() + for path in "${RUNNER_TEMP}/fail-closed-budget-e2e.log" \ + "${RUNNER_TEMP}/fail-closed-budget-e2e.xml" "${RUNNER_TEMP}/fail-closed-budget-proxy.log"; do + if [ -f "${path}" ]; then + files+=("${path}") + fi + done + if [ "${#files[@]}" = "0" ]; then + exit 0 + fi + uv run --no-sync python .github/e2e-stack/redact_output.py \ + --values tests/e2e/.env --values "${RUNNER_TEMP}/fail-closed-budget-values.env" \ + --out "${RUNNER_TEMP}/fail-closed-budget-redacted" "${files[@]}" + + - name: Keep redacted e2e output + if: always() && steps.proxy.outcome != 'skipped' + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: fail-closed-team-member-budget-e2e-${{ github.run_attempt }} + path: ${{ runner.temp }}/fail-closed-budget-redacted + retention-days: 14 + if-no-files-found: ignore + + - name: Stop proxy and remove credentials + if: always() + run: | + if [ -n "${E2E_PROXY_PID:-}" ]; then + kill "${E2E_PROXY_PID}" 2>/dev/null || true + fi + rm -f tests/e2e/.env "${RUNNER_TEMP}/fail-closed-budget-values.env" \ + "${RUNNER_TEMP}/fail-closed-budget-e2e.log" "${RUNNER_TEMP}/fail-closed-budget-e2e.xml" \ + "${RUNNER_TEMP}/fail-closed-budget-proxy.log" + rm -rf "${RUNNER_TEMP}/fail-closed-budget-redacted" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3ec430332ee..08ac80d9495 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5885,6 +5885,7 @@ async def _check_team_member_budget( counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", fallback_spend=team_member_spend, max_budget=team_member_budget, + fallback_authoritative=loaded_membership is None, ) if not math.isfinite(team_member_budget): diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 603591006d9..75752aa4b0d 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -66,6 +66,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "fail_closed_budget_stack": "E2E_FAIL_CLOSED_BUDGET_STACK", "cli_determinism": CLI_DETERMINISM_OPT_IN_ENV, "mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV, "provider_edge_host": PROVIDER_EDGE_HOST_OPT_IN_ENV, diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 163de67fc41..af26f518941 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -21,6 +21,7 @@ - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} - {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} +- {id: quota_management.budget.team_member.missing_membership_counts_as_verified_zero_spend, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [missing_membership_counts_as_verified_zero_spend], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py + proxy_server.py:get_current_spend", rationale: "A key scoped to a team with a default member budget but no membership row treats the absent row as verified zero spend and admits the request"} - {id: quota_management.budget.model_access_group.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A model access group's shared max_budget blocks further calls to deployments in the group once the pool is spent"} - {id: quota_management.budget.model_access_group.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "The pool is shared, so a key that spent nothing of its own is blocked once another key granted the same group drained it"} - {id: quota_management.budget.model_access_group.isolates_per_group, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [isolates_per_group], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A request is charged only to the granted groups that serve the model it called, so an exhausted group never blocks a sibling group"} diff --git a/tests/e2e/gateway/fail_closed_team_member_budget_ci_config.yml b/tests/e2e/gateway/fail_closed_team_member_budget_ci_config.yml new file mode 100644 index 00000000000..f7f45716d69 --- /dev/null +++ b/tests/e2e/gateway/fail_closed_team_member_budget_ci_config.yml @@ -0,0 +1,19 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true + fail_closed_budget_enforcement: true + +litellm_settings: + enable_redis_auth_cache: true + cache: true + cache_params: + type: redis + host: 127.0.0.1 + port: 6379 + socket_timeout: 0.1 + +model_list: + - model_name: claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + api_key: os.environ/ANTHROPIC_API_KEY diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index d01caeff3ea..eee83c8ea7e 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -12,6 +12,7 @@ markers = prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + fail_closed_budget_stack: needs a proxy booted from gateway/fail_closed_team_member_budget_ci_config.yml; deselected unless E2E_FAIL_CLOSED_BUDGET_STACK is set quiet_stack: measures the proxy itself, so it runs while no other test on this host is hitting the stack; every other test waits for it to finish mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set provider_edge_host: routes provider traffic through the pytest host's edge in every fixture mode, so the gateway must reach the pytest host; deselected unless E2E_PROVIDER_EDGE_HOST_REACHABLE is set diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 087dc8ca522..78729bc8e5c 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -97,6 +97,7 @@ class TeamNewBody(BaseModel): max_budget: float | None = None budget_duration: str | None = None organization_id: str | None = None + team_member_budget: float | None = None budget_limits: list[BudgetWindow] | None = None @@ -416,6 +417,7 @@ class BudgetClient: max_budget: float | None = None, budget_duration: str | None = None, organization_id: str | None = None, + team_member_budget: float | None = None, budget_limits: list[BudgetWindow] | None = None, ) -> str: team_id = unwrap( @@ -427,6 +429,7 @@ class BudgetClient: max_budget=max_budget, budget_duration=budget_duration, organization_id=organization_id, + team_member_budget=team_member_budget, budget_limits=budget_limits, ), response_type=TeamNewResponse, diff --git a/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py b/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py index 0fd0a545660..5563bdc1a7e 100644 --- a/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py +++ b/tests/e2e/quota_management/budgets/test_team_member_budget_e2e.py @@ -61,6 +61,25 @@ def member(client: BudgetClient) -> Iterator[_Member]: resources.teardown() +@pytest.fixture(scope="class") +def member_without_membership(client: BudgetClient) -> Iterator[_Member]: + resources = ResourceManager(client=client.proxy) + try: + marker = unique_marker() + team_id = client.create_team( + alias=f"e2e-team-member-no-membership-{marker}", + team_member_budget=TEAM_BUDGET, + ) + resources.defer(lambda: client.delete_team(team_id)) + user_id = client.create_user(max_budget=TEAM_BUDGET) + resources.defer(lambda: client.delete_user(user_id)) + key = client.generate_key(team_id=team_id, user_id=user_id) + resources.defer(lambda: client.delete_key(key)) + yield _Member(team_id=team_id, user_id=user_id, key=key) + finally: + resources.teardown() + + def _send(client: BudgetClient, key: str) -> str | None: """One member call; its response id (== the spend-log request_id) if it went through, else None.""" @@ -83,9 +102,7 @@ class TestTeamMemberBudget: sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid) assert sent, "no member call went through; cannot check attribution" - rows = client.proxy.poll_logs_for_key( - member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs}) - ) + rows = client.proxy.poll_logs_for_key(member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs})) logged = [row for row in rows if row.request_id in sent] assert logged, f"none of the member's {len(sent)} calls reached the spend logs" @@ -106,3 +123,18 @@ class TestTeamMemberBudget: require_successful_call(result) time.sleep(2) pytest.fail("per-member budget never enforced within the call budget") + + +class TestFailClosedTeamMemberBudgetWithoutMembership: + @pytest.mark.fail_closed_budget_stack + @pytest.mark.covers("quota_management.budget.team_member.missing_membership_counts_as_verified_zero_spend") + def test_missing_membership_counts_as_verified_zero_spend( + self, client: BudgetClient, member_without_membership: _Member + ) -> None: + result = client.chat( + member_without_membership.key, + MODEL, + f"missing membership {unique_marker()}", + max_tokens=16, + ) + assert result.status_code == 200, result.body diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 353249dddf0..c34afaaf3c7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6,7 +6,7 @@ import sys import time from collections.abc import Iterator, Mapping from types import SimpleNamespace -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch if TYPE_CHECKING: @@ -17,9 +17,10 @@ from datetime import datetime, timedelta, timezone import httpx import pytest -from fastapi import Request, status +from fastapi import HTTPException, Request, status import litellm +from litellm.proxy import proxy_server from litellm.proxy._types import ( CallInfo, Litellm_EntityType, @@ -28,6 +29,7 @@ from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LiteLLM_TagTable, LiteLLM_TeamTable, + LiteLLM_TeamMembership, LiteLLM_UserTable, LitellmUserRoles, ModelAccessDeniedProxyException, @@ -68,6 +70,7 @@ from litellm.proxy.auth.auth_checks import ( route_skips_budget_checks, vector_store_access_check, ) +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -83,6 +86,7 @@ from litellm.proxy.auth.user_api_key_auth import check_api_key_for_custom_header from litellm.proxy import proxy_server from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_bearer_token, encrypt_value_helper from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.utils import PrismaClient, ProxyLogging from prisma.errors import DataError from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, @@ -8086,6 +8090,96 @@ async def test_check_team_member_budget_fails_closed_when_the_membership_read_hi await _check_team_member_budget(user_object=None, **_restricted_member_check_deps()) +def _unavailable_spend_counter_cache() -> DualCache: + redis_cache: Final = cast( + RedisCache, + MagicMock(async_get_cache=AsyncMock(side_effect=RuntimeError("redis unavailable"))), + ) + return DualCache(redis_cache=redis_cache) + + +def _prisma_client_with_membership_lookup(find_unique: AsyncMock) -> PrismaClient: + return cast( + PrismaClient, + SimpleNamespace( + db=SimpleNamespace(litellm_teammembership=SimpleNamespace(find_unique=find_unique)), + ), + ) + + +@pytest.mark.asyncio +async def test_check_team_member_budget_missing_membership_is_verified_with_unavailable_counters(): + cache: Final = cast( + UserApiKeyCache, + MagicMock( + async_get_cache=AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="default-budget-100", max_budget=100.0) + ) + ), + ) + membership_find_unique: Final = AsyncMock(return_value=None) + prisma_client: Final = _prisma_client_with_membership_lookup(membership_find_unique) + + with ( + patch.object(proxy_server, "general_settings", {"fail_closed_budget_enforcement": True}), + patch.object(proxy_server, "prisma_client", prisma_client), + patch.object(proxy_server, "spend_counter_cache", _unavailable_spend_counter_cache()), + ): + result: Final = await _check_team_member_budget( + team_object=LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "default-budget-100"}, + ), + user_object=None, + valid_token=UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team"), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + team_membership=None, + team_membership_loaded=True, + ) + + assert result is None + membership_find_unique.assert_awaited() + + +@pytest.mark.asyncio +async def test_check_team_member_budget_existing_membership_still_fails_closed_when_counters_are_unavailable(): + cache: Final = cast( + UserApiKeyCache, + MagicMock( + async_get_cache=AsyncMock( + return_value=LiteLLM_BudgetTable(budget_id="default-budget-100", max_budget=100.0) + ) + ), + ) + membership_find_unique: Final = AsyncMock(side_effect=RuntimeError("database unavailable")) + prisma_client: Final = _prisma_client_with_membership_lookup(membership_find_unique) + + with ( + patch.object(proxy_server, "general_settings", {"fail_closed_budget_enforcement": True}), + patch.object(proxy_server, "prisma_client", prisma_client), + patch.object(proxy_server, "spend_counter_cache", _unavailable_spend_counter_cache()), + pytest.raises(HTTPException) as exc_info, + ): + await _check_team_member_budget( + team_object=LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "default-budget-100"}, + ), + user_object=None, + valid_token=UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team"), + prisma_client=prisma_client, + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + team_membership=LiteLLM_TeamMembership(user_id="test-user", team_id="test-team", spend=0.0), + team_membership_loaded=True, + ) + + assert exc_info.value.status_code == 503 + membership_find_unique.assert_awaited() + + @pytest.mark.asyncio async def test_get_team_membership_waiter_cancel_does_not_cancel_shared_load(): from litellm.proxy.auth.auth_checks import get_team_membership diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index 922504ecc58..64f9b3b1540 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -884,7 +884,14 @@ async def test_team_update_reaches_inherited_members_but_not_overridden_ones(): team_member_budget=1.0, ) - async def spend_from_membership(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + async def spend_from_membership( + counter_key: str, + fallback_spend: float, + max_budget: float | None = None, + *, + fallback_authoritative: bool, + ) -> float: + assert not fallback_authoritative return fallback_spend async def check(user_id: str, spend: float) -> None: