Merge ab004e5ae3 into f285229b51

.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"

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):

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,

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"}

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

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

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,

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

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

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: