Merge pull request #26356 from BerriAI/litellm_cci_gha_dedup_and_shard

[Infra] Remove CCI/GHA test duplication and semantically shard proxy DB tests

.circleci/config.yml

@@ -576,39 +576,6 @@ jobs:
           paths:
             - realtime_translation_coverage.xml
             - realtime_translation_coverage
-  mcp_testing:
-    docker:
-      - *python312_image
-    working_directory: ~/project
-
-    steps:
-      - checkout
-      - setup_google_dns
-      - install_uv
-      - run:
-          name: Install Dependencies
-          command: |
-            uv sync --frozen --all-groups --all-extras --python 3.12
-      # Run pytest and generate JUnit XML report
-      - run:
-          name: Run tests
-          command: |
-            uv run --no-sync python -m pytest -vv tests/mcp_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 -n 2
-          no_output_timeout: 15m
-      - run:
-          name: Rename the coverage files
-          command: |
-            mv coverage.xml mcp_coverage.xml
-            mv .coverage mcp_coverage
-
-      # Store test results
-      - store_test_results:
-          path: test-results
-      - persist_to_workspace:
-          root: .
-          paths:
-            - mcp_coverage.xml
-            - mcp_coverage
   agent_testing:
     docker:
       - *python312_image
@@ -807,39 +774,6 @@ jobs:
           paths:
             - search_coverage.xml
             - search_coverage
-  # Split litellm_mapped_tests into parallel jobs
-  litellm_mapped_tests_proxy_part1:
-    docker:
-      - *python312_image
-    working_directory: ~/project
-    resource_class: large
-    steps:
-      - setup_litellm_test_deps
-      - run:
-          name: Run proxy tests part 1 (high-volume directories)
-          command: |
-            uv run --no-sync python -m prisma generate
-            export PYTHONUNBUFFERED=1
-            uv run --no-sync python -m pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/client tests/test_litellm/proxy/auth --junitxml=test-results/junit-proxy-part1.xml --durations=10 -n 4 --maxfail=5 --timeout=60 -vv --log-cli-level=WARNING -r A
-          no_output_timeout: 15m
-      - store_test_results:
-          path: test-results
-  litellm_mapped_tests_proxy_part2:
-    docker:
-      - *python312_image
-    working_directory: ~/project
-    resource_class: large
-    steps:
-      - setup_litellm_test_deps
-      - run:
-          name: Run proxy tests part 2 (all other tests)
-          command: |
-            uv run --no-sync python -m prisma generate
-            export PYTHONUNBUFFERED=1
-            uv run --no-sync python -m pytest tests/test_litellm/proxy --ignore=tests/test_litellm/proxy/guardrails --ignore=tests/test_litellm/proxy/management_endpoints --ignore=tests/test_litellm/proxy/_experimental --ignore=tests/test_litellm/proxy/client --ignore=tests/test_litellm/proxy/auth --junitxml=test-results/junit-proxy-part2.xml --durations=10 -n 4 --maxfail=5 --timeout=120 -vv --log-cli-level=WARNING -r A
-          no_output_timeout: 15m
-      - store_test_results:
-          path: test-results
   litellm_mapped_enterprise_tests:
     docker:
       - *python312_image
@@ -2091,7 +2025,7 @@ jobs:
       - run:
           name: Combine Coverage
           command: |
-            uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage mcp_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage redis_caching_coverage
+            uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage redis_caching_coverage
             uv tool run --from 'coverage[toml]==7.10.6' coverage xml
       - codecov/upload:
           file: ./coverage.xml
@@ -2426,8 +2360,6 @@ workflows:
           filters: *main_branches
       - realtime_translation_testing:
           filters: *main_branches
-      - mcp_testing:
-          filters: *main_branches
       - agent_testing:
           filters: *main_branches
       - guardrails_testing:
@@ -2442,10 +2374,6 @@ workflows:
           filters: *main_branches
       - litellm_mapped_enterprise_tests:
           filters: *main_branches
-      - litellm_mapped_tests_proxy_part1:
-          filters: *main_branches
-      - litellm_mapped_tests_proxy_part2:
-          filters: *main_branches
       - batches_testing:
           filters: *main_branches
       - litellm_utils_testing:
@@ -2463,14 +2391,11 @@ workflows:
       - upload-coverage:
           requires:
             - realtime_translation_testing
-            - mcp_testing
             - agent_testing
             - google_generate_content_endpoint_testing
             - guardrails_testing
             - ocr_testing
             - search_testing
-            - litellm_mapped_tests_proxy_part1
-            - litellm_mapped_tests_proxy_part2
             - litellm_mapped_enterprise_tests
             - batches_testing
             - litellm_utils_testing

.github/workflows/_test-unit-services-base.yml

@@ -32,41 +32,39 @@ on:
         required: false
         type: boolean
         default: false
+      dist:
+        description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
+        required: false
+        type: string
+        default: "loadscope"
       artifact-name:
         description: "Unique name for the coverage artifact (must be unique per run)"
         required: false
         type: string
         default: "run"
-    secrets:
-      DATABASE_URL:
-        required: false
-      POSTGRES_USER:
-        required: false
-      POSTGRES_PASSWORD:
-        required: false
 
 permissions:
   contents: read
 
+# The postgres service container below is spawned per-job on localhost and
+# destroyed with the job. Nothing outside the runner can reach it. The
+# user/password/database here are not secrets — they're bootstrap values
+# for a throwaway container — so we hardcode them instead of attaching
+# every matrix shard to a GHA environment just to read three "secrets"
+# (which also produces a "temporarily deployed to …" notification on the
+# PR timeline per shard per push).
 jobs:
   run:
     name: Run tests
     runs-on: ubuntu-latest
     timeout-minutes: ${{ inputs.timeout-minutes }}
-    # Environment is derived from the enable-* flags, not caller-controllable.
-    # This prevents callers from passing arbitrary environment names to bypass secret scoping.
-    environment: >-
-      ${{
-        inputs.enable-postgres && 'integration-postgres' ||
-        ''
-      }}
 
     services:
       postgres:
         image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
         env:
-          POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
-          POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+          POSTGRES_USER: litellm
+          POSTGRES_PASSWORD: litellm
           POSTGRES_DB: litellm_test
         ports:
           - 5432:5432
@@ -114,7 +112,7 @@ jobs:
       - name: Run Prisma migrations
         if: ${{ inputs.enable-postgres }}
         env:
-          DATABASE_URL: ${{ secrets.DATABASE_URL }}
+          DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
         run: |
           uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
 
@@ -124,7 +122,8 @@ jobs:
           MAX_FAILURES: ${{ inputs.max-failures }}
           WORKERS: ${{ inputs.workers }}
           RERUNS: ${{ inputs.reruns }}
-          DATABASE_URL: ${{ inputs.enable-postgres && secrets.DATABASE_URL || '' }}
+          DIST: ${{ inputs.dist }}
+          DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
         run: |
           if [ "${WORKERS}" = "0" ]; then
             uv run --no-sync pytest ${TEST_PATH:?} \
@@ -143,7 +142,7 @@ jobs:
               -n "${WORKERS}" \
               --reruns "${RERUNS}" \
               --reruns-delay 1 \
-              --dist=loadscope \
+              --dist="${DIST}" \
               --durations=20 \
               --cov=litellm \
               --cov-report=xml:coverage.xml \

.github/workflows/test-unit-proxy-db.yml

@@ -12,8 +12,74 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: true
 
+# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
+# rather than alphabetical letter ranges. Adding a new test file means adding it
+# to whichever group it belongs to, not reshuffling slices.
+#
+# Design targets:
+#   * Every shard runs in <= 7 minutes of wall-clock on the default runner.
+#     Most of a shard's time is pytest plugin load + xdist worker imports +
+#     pytest-cov instrumentation, not the tests themselves. Keeping per-shard
+#     work low and matching worker count to runner cores is what controls it.
+#   * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores
+#     oversubscribes 2x and workers fight for CPU during their cold-start
+#     imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective).
+#   * test_key_generate_prisma.py stays serial (workers=0) — it has event-loop
+#     conflicts with the logging worker when run in parallel.
+#   * test_proxy_utils.py runs as a single shard with --dist=worksteal so
+#     xdist balances its 188 parametrized cases across workers instead of
+#     pinning the whole file to one worker (the default --dist=loadscope
+#     behavior for single-file targets).
+#   * test_db_schema_migration.py is isolated because one test in it
+#     (test_aaaasschema_migration_check) takes ~170s — by itself it
+#     determines the shard's wall-clock floor.
 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, 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 it belongs to.")
+              sys.exit(1)
+          print(f"OK: all {len(actual)} files assigned to a shard.")
+          PY
+
   proxy-db:
+    needs: assert-shard-coverage
+    # Display only the semantic shard name in the checks UI instead of GHA's
+    # default "proxy-db (key-generation, tests/proxy_unit_tests/…, 0, loadscope, 20)"
+    # which includes every matrix field and gets truncated past the test-path.
+    name: ${{ matrix.test-group }}
     permissions:
       contents: read
       id-token: write
@@ -22,26 +88,146 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          # Key generation tests must NOT run in parallel (event loop conflicts with logging worker)
+          # Must run serially — event-loop conflict with the logging worker.
           - test-group: key-generation
             test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py"
             workers: 0
-            timeout: 30
-          - test-group: auth-checks
-            test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py"
-            workers: 8
+            dist: loadscope
             timeout: 20
-          # test_proxy_utils.py is large (168+ parametrized tests) — run it on its
-          # own matrix so --dist=loadscope doesn't pin all of it to a single xdist
-          # worker and push the "remaining" group past the job timeout.
+
+          # ---- auth: split into 2 shards ----
+          - test-group: auth-checks
+            test-path: >-
+              tests/proxy_unit_tests/test_auth_checks.py
+              tests/proxy_unit_tests/test_user_api_key_auth.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+          - test-group: jwt-and-keys
+            test-path: >-
+              tests/proxy_unit_tests/test_jwt.py
+              tests/proxy_unit_tests/test_jwt_key_mapping.py
+              tests/proxy_unit_tests/test_proxy_custom_auth.py
+              tests/proxy_unit_tests/test_key_generate_dynamodb.py
+              tests/proxy_unit_tests/test_deployed_proxy_keygen.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+
+          # ---- test_proxy_utils.py, single shard, worksteal distribution ----
           - test-group: proxy-utils
             test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
-            workers: 8
-            timeout: 20
-          - test-group: remaining
-            test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --ignore=tests/proxy_unit_tests/test_proxy_utils.py"
-            workers: 8
-            timeout: 30
+            workers: 4
+            dist: worksteal
+            timeout: 15
+
+          # ---- proxy server: split into 2 shards ----
+          - test-group: proxy-server-core
+            test-path: >-
+              tests/proxy_unit_tests/test_proxy_server.py
+              tests/proxy_unit_tests/test_proxy_server_keys.py
+              tests/proxy_unit_tests/test_proxy_server_caching.py
+              tests/proxy_unit_tests/test_proxy_server_langfuse.py
+              tests/proxy_unit_tests/test_proxy_server_spend.py
+              tests/proxy_unit_tests/test_aproxy_startup.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+          - test-group: proxy-runtime
+            test-path: >-
+              tests/proxy_unit_tests/test_proxy_config_unit_test.py
+              tests/proxy_unit_tests/test_proxy_routes.py
+              tests/proxy_unit_tests/test_proxy_gunicorn.py
+              tests/proxy_unit_tests/test_server_root_path.py
+              tests/proxy_unit_tests/test_proxy_pass_user_config.py
+              tests/proxy_unit_tests/test_proxy_token_counter.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+
+          # ---- logging: split into 2 shards ----
+          - test-group: custom-logging
+            test-path: >-
+              tests/proxy_unit_tests/test_custom_callback_input.py
+              tests/proxy_unit_tests/test_custom_logger_s3_gcs.py
+              tests/proxy_unit_tests/test_proxy_custom_logger.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+          - test-group: logging-misc
+            test-path: >-
+              tests/proxy_unit_tests/test_proxy_reject_logging.py
+              tests/proxy_unit_tests/test_audit_logs_proxy.py
+              tests/proxy_unit_tests/test_search_api_logging.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+
+          # ---- db-and-spend: isolate the 170s schema-migration test ----
+          # test_db_schema_migration.py has exactly one test, and that test
+          # is mostly waiting on `prisma migrate deploy` / `prisma migrate
+          # diff` subprocesses (~170s). It does no CPU-bound Python work
+          # inside the test. Running with workers=0 (serial, no xdist)
+          # skips the 4-worker cold-start cost we'd otherwise pay for a
+          # single test, saving ~4 minutes of wall-clock.
+          - test-group: schema-migration
+            test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
+            workers: 0
+            dist: loadscope
+            timeout: 15
+          - test-group: db-and-spend
+            test-path: >-
+              tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
+              tests/proxy_unit_tests/test_db_schema_changes.py
+              tests/proxy_unit_tests/test_e2e_pod_lock_manager.py
+              tests/proxy_unit_tests/test_skills_db.py
+              tests/proxy_unit_tests/test_update_daily_tag_spend.py
+              tests/proxy_unit_tests/test_update_spend.py
+              tests/proxy_unit_tests/test_project_endpoints_prisma.py
+              tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+
+          # ---- guardrails + budget + hooks: split into 2 ----
+          - test-group: guardrails-hooks
+            test-path: >-
+              tests/proxy_unit_tests/test_proxy_setting_guardrails.py
+              tests/proxy_unit_tests/test_banned_keyword_list.py
+              tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+          - test-group: budgets
+            test-path: >-
+              tests/proxy_unit_tests/test_default_end_user_budget_simple.py
+              tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
+              tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
+            workers: 4
+            dist: loadscope
+            timeout: 15
+
+          - test-group: endpoints-and-responses
+            test-path: >-
+              tests/proxy_unit_tests/test_blog_posts_endpoint.py
+              tests/proxy_unit_tests/test_models_fallback_endpoint.py
+              tests/proxy_unit_tests/test_google_endpoint_routing.py
+              tests/proxy_unit_tests/test_google_gemini_proxy_request.py
+              tests/proxy_unit_tests/test_get_favicon.py
+              tests/proxy_unit_tests/test_get_image.py
+              tests/proxy_unit_tests/test_ui_path_detection.py
+              tests/proxy_unit_tests/test_prompt_test_endpoint.py
+              tests/proxy_unit_tests/test_check_batch_cost.py
+              tests/proxy_unit_tests/test_check_responses_cost.py
+              tests/proxy_unit_tests/test_response_polling_handler.py
+              tests/proxy_unit_tests/test_response_polling_pre_call_checks.py
+              tests/proxy_unit_tests/test_realtime_cache.py
+              tests/proxy_unit_tests/test_proxy_exception_mapping.py
+              tests/proxy_unit_tests/test_custom_tokenizer_bug.py
+              tests/proxy_unit_tests/test_model_response_typing
+            workers: 4
+            dist: loadscope
+            timeout: 15
     uses: ./.github/workflows/_test-unit-services-base.yml
     with:
       test-path: ${{ matrix.test-path }}
@@ -49,8 +235,5 @@ jobs:
       reruns: 2
       timeout-minutes: ${{ matrix.timeout }}
       enable-postgres: true
+      dist: ${{ matrix.dist }}
       artifact-name: proxy-db-${{ matrix.test-group }}
-    secrets:
-      DATABASE_URL: ${{ secrets.DATABASE_URL }}
-      POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
-      POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}

.github/workflows/test-unit-proxy-endpoints.yml

@@ -36,6 +36,8 @@ jobs:
         tests/test_litellm/proxy/health_endpoints
         tests/test_litellm/proxy/public_endpoints
         tests/test_litellm/proxy/prompts
+        tests/test_litellm/proxy/rag_endpoints
+        tests/test_litellm/proxy/realtime_endpoints
         tests/test_litellm/proxy/ui_crud_endpoints
       workers: 2
       reruns: 2

.github/workflows/test-unit-security.yml

@@ -1,6 +1,8 @@
 name: "Unit Tests: Security"
 
-# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
+# Kept push-only (was previously required by DATABASE_URL secret scoping;
+# now the postgres credentials are ephemeral localhost values but the
+# push-trigger stays to match the proxy-db workflow cadence).
 on:
   push:
     branches: [main, "litellm_**"]
@@ -24,7 +26,3 @@ jobs:
       timeout-minutes: 20
       enable-postgres: true
       artifact-name: security
-    secrets:
-      DATABASE_URL: ${{ secrets.DATABASE_URL }}
-      POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
-      POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}

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