From 32c390a0f6afd3acc85aed4492aa154d6a2d401d Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 23 Apr 2026 15:01:25 -0700
Subject: [PATCH] fix(tests): restore proxy_server.master_key in realtime
 fixture; add shard-coverage guard

Two fixes to proxy-db CI:

1. test_realtime_webrtc_endpoints.py's `proxy_app` fixture mutated the
   module-global `proxy_server.master_key` without restoring it, leaking
   state into any test that shared the same xdist worker. Under
   --dist=loadscope with 2 workers (GHA proxy-endpoints), this caused the
   google_endpoints tests to fail with "No api key passed in." because
   user_api_key_auth saw a set master_key and a missing API key on the
   test request. The fixture now saves and restores the original value.

2. Address the Greptile note that the semantic shard design has no
   catch-all, so a new test file added to tests/proxy_unit_tests/ without
   a matrix entry would silently skip CI. Adds an assert-shard-coverage
   job that enumerates test_*.py files and fails the workflow if any are
   not referenced by a matrix entry, with a clear message telling the
   author which semantic shard to place it in. All proxy-db shards now
   depend on this guard.
---
 .github/workflows/test-unit-proxy-db.yml      | 43 +++++++++++++++++++
 .../test_realtime_webrtc_endpoints.py         | 10 ++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml
index f8d5bc265a..0f2694984e 100644
--- a/.github/workflows/test-unit-proxy-db.yml
+++ b/.github/workflows/test-unit-proxy-db.yml
@@ -25,7 +25,50 @@ concurrency:
 #     pinning the whole file to a single worker (the default --dist=loadscope
 #     behavior for single-file targets).
 jobs:
+  # Fast guard — fails the workflow if a test_*.py file under
+  # tests/proxy_unit_tests/ is not referenced by any matrix entry below.
+  # The semantic-shard design (no catch-all "remaining" bucket) relies on
+  # every test file being explicitly assigned; this guard prevents a new
+  # file from silently dropping out of CI.
+  assert-shard-coverage:
+    runs-on: ubuntu-latest
+    timeout-minutes: 2
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          persist-credentials: false
+      - name: Assert every test_*.py is in a matrix shard
+        run: |
+          python3 - <<'PY'
+          import pathlib, re, sys, yaml
+          wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
+          matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
+          referenced = set()
+          for entry in matrix:
+              for token in entry["test-path"].split():
+                  if token.startswith("tests/proxy_unit_tests/"):
+                      referenced.add(pathlib.PurePosixPath(token).name)
+          actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
+                    if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
+                    and p.name != "test_configs"}
+          orphans = sorted(actual - referenced)
+          if orphans:
+              print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
+              print("       are not assigned to any shard in test-unit-proxy-db.yml:")
+              for o in orphans:
+                  print(f"  - {o}")
+              print()
+              print("Add each to whichever semantic shard (auth-and-jwt, proxy-server,")
+              print("logging-and-callbacks, db-and-spend, guardrails-budget-hooks,")
+              print("endpoints-and-responses, proxy-utils, key-generation) it belongs to.")
+              sys.exit(1)
+          print(f"OK: all {len(actual)} files assigned to a shard.")
+          PY
+
   proxy-db:
+    needs: assert-shard-coverage
     permissions:
       contents: read
       id-token: write
diff --git a/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py b/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py
index e414f975f5..99b6335ce8 100644
--- a/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py
+++ b/tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py
@@ -116,8 +116,16 @@ def test_decode_realtime_token_payload_ephemeral_key_not_string():
 def proxy_app():
     from litellm.proxy import proxy_server
 
+    # master_key is a module-global — restore it on teardown so this fixture
+    # doesn't leak state into unrelated tests that share the same xdist worker
+    # (e.g. tests that assume master_key is None and send unauthenticated
+    # requests to the shared FastAPI app).
+    original_master_key = proxy_server.master_key
     proxy_server.master_key = "sk-test-master-key"
-    return proxy_server.app
+    try:
+        yield proxy_server.app
+    finally:
+        proxy_server.master_key = original_master_key
 
 
 @pytest.fixture