ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests (#42902)

* ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure

* test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage

* test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split

* ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags

* ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags

* ci: fail the unit shard when circleci tests split errors

* test: drop restating comments from the gemini pass-through split

* ci: exit the unit shard cleanly when circleci tests split assigns it no files

---------

Co-authored-by: yuneng <yuneng@berri.ai>

.circleci/scripts/unit_selection.sh

@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+flag="${1:?usage: unit_selection.sh <codecov flag>}"
+
+legacy_flags=(
+  caching-local
+  enterprise-package
+  enterprise-routing
+  proxy-extras
+  proxy-infra
+)
+
+legacy_paths() {
+  case "$1" in
+    caching-local) echo tests/unit/caching ;;
+    enterprise-package)
+      echo tests/unit/enterprise/integrations
+      echo tests/unit/enterprise/proxy/auth
+      echo tests/unit/enterprise/proxy/guardrails
+      echo tests/unit/enterprise/proxy/hooks
+      echo tests/unit/enterprise/proxy/management_endpoints
+      echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py
+      echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;;
+    enterprise-routing)
+      echo tests/unit/enterprise/enterprise_callbacks/send_emails
+      echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
+      echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py
+      echo tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py
+      echo tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py
+      echo tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py
+      echo tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py
+      echo tests/unit/enterprise/proxy/test_enterprise_routes.py
+      echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py
+      echo tests/unit/enterprise/proxy/test_managed_files_access_check.py
+      echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;;
+    proxy-extras) echo tests/unit/litellm_proxy_extras ;;
+    proxy-infra) echo tests/unit/gateway ;;
+    *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;;
+  esac
+}
+
+expand() {
+  while read -r path; do
+    if [ -d "$path" ]; then
+      find "$path" -name 'test_*.py'
+    elif [ -f "$path" ]; then
+      echo "$path"
+    else
+      echo "unit_selection.sh: $path does not exist" >&2
+      exit 1
+    fi
+  done
+}
+
+if [ "$flag" = unit ]; then
+  comm -23 \
+    <(find tests/unit -name 'test_*.py' | sort) \
+    <(for legacy in "${legacy_flags[@]}"; do legacy_paths "$legacy"; done | expand | sort)
+  exit 0
+fi
+
+legacy_paths "$flag" | expand | sort

.circleci/tests.yml

@@ -171,6 +171,12 @@ jobs:
       shards:
         type: integer
         default: 6
+      workers:
+        type: integer
+        default: 4
+      dist:
+        type: string
+        default: loadscope
       base_ref:
         type: string
         default: ""
@@ -199,17 +205,19 @@ jobs:
           no_output_timeout: 20m
           command: |
             mkdir -p test-results/<< parameters.flag >>
-            selection="$(find tests/unit -name 'test_*.py' | sort)" || { echo "test selection failed for << parameters.flag >>"; exit 1; }
-            [ -n "${selection}" ] || { echo "test selection produced no files for << parameters.flag >>"; exit 1; }
+            selection="$(bash .circleci/scripts/unit_selection.sh << parameters.flag >>)" || { echo "unit_selection.sh failed for << parameters.flag >>"; exit 1; }
+            [ -n "${selection}" ] || { echo "unit_selection.sh produced no files for << parameters.flag >>"; exit 1; }
             shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; }
             [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; }
             mapfile -t files < <(printf '%s\n' "${shard}")
+            xdist_args=()
+            if [ "<< parameters.workers >>" -gt 0 ]; then xdist_args=(-n << parameters.workers >> --dist=<< parameters.dist >>); fi
             rerun_args=(-p no:rerunfailures)
             if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi
             test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP")
             set +e
             env -i "${test_env[@]}" \
-              uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml
+              uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml
             status=$?
             set -e
             if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi
@@ -293,6 +301,25 @@ workflows:
       - unit:
           base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
           pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
+      - unit:
+          name: unit-<< matrix.flag >>
+          shards: 1
+          workers: 2
+          reruns: 2
+          matrix:
+            parameters:
+              flag: [caching-local, proxy-extras, enterprise-routing]
+          base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
+          pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
+      - unit:
+          name: unit-<< matrix.flag >>
+          shards: 1
+          reruns: 2
+          matrix:
+            parameters:
+              flag: [enterprise-package, proxy-infra]
+          base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >>
+          pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >>
       - documentation
       - integration:
           name: integration-<< matrix.suite >>

.github/scripts/assert_ci_coverage.py

@@ -120,6 +120,13 @@ def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
     )
 
 
+def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
+    script: Final = repo_root / ".circleci/scripts/unit_selection.sh"
+    if not script.is_file():
+        return frozenset()
+    return frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(_uncommented(script.read_text())))
+
+
 def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
     return frozenset(
         match.group(0)
@@ -611,7 +618,10 @@ def main() -> int:
     scalars = _all_scalars()
 
     integration_paths, ownership_findings = _integration_ownership()
-    test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | integration_paths) + ownership_findings
+    test_findings = (
+        _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | _unit_selection_tokens() | integration_paths)
+        + ownership_findings
+    )
     dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
     stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
 

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

@@ -13,6 +13,15 @@ on:
           have its path existence-checked like any other token.
         required: true
         type: string
+      fork-flag:
+        description: >-
+          Codecov flag of the `.circleci/tests.yml` job that now owns part of
+          this shard. CircleCI does not run on pull requests from forks, so on
+          those events this shard also runs the files
+          `.circleci/scripts/unit_selection.sh` lists for the flag.
+        required: false
+        type: string
+        default: ""
       workers:
         description: "Number of pytest-xdist workers"
         required: false
@@ -92,6 +101,7 @@ jobs:
       pull-requests: read
     outputs:
       decision: ${{ steps.changes.outputs.decision }}
+      has-coverage: ${{ steps.tests.outputs.has-coverage }}
 
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@@ -160,10 +170,13 @@ jobs:
           uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
 
       - name: Run tests
+        id: tests
         if: steps.changes.outputs.decision != 'skip'
         timeout-minutes: ${{ inputs.timeout-minutes }}
         env:
           TEST_PATH: ${{ inputs.test-path }}
+          FORK_FLAG: ${{ inputs.fork-flag }}
+          IS_FORK: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository }}
           MAX_FAILURES: ${{ inputs.max-failures }}
           WORKERS: ${{ inputs.workers }}
           RERUNS: ${{ inputs.reruns }}
@@ -171,9 +184,18 @@ jobs:
           DIST: ${{ inputs.dist }}
           COVERAGE_CORE: sysmon
         run: |
+          echo "has-coverage=false" >> "$GITHUB_OUTPUT"
+          selection="${TEST_PATH}"
+          if [ "${IS_FORK}" = "true" ] && [ -n "${FORK_FLAG}" ]; then
+            selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${FORK_FLAG}" | tr '\n' ' ')"
+          fi
+          if [ -z "${selection// /}" ]; then
+            echo "shard selection is empty on this event (CircleCI flag ${FORK_FLAG:-none} owns it); nothing to run"
+            exit 0
+          fi
           pytest_args=()
           existing_paths=0
-          for token in ${TEST_PATH:?}; do
+          for token in ${selection}; do
             case "${token}" in
               -*) pytest_args+=("${token}") ;;
               *)
@@ -187,7 +209,7 @@ jobs:
             esac
           done
           if [ "${existing_paths}" -eq 0 ]; then
-            echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run"
+            echo "No path in the selection exists (${selection}); nothing to run"
             exit 0
           fi
           xdist_args=()
@@ -209,8 +231,11 @@ jobs:
             --cov-config=pyproject.toml
           status=$?
           set -e
+          if [ -f coverage.xml ]; then
+            echo "has-coverage=true" >> "$GITHUB_OUTPUT"
+          fi
           if [ "$status" -eq 5 ]; then
-            echo "pytest collected no tests from ${TEST_PATH}; passing"
+            echo "pytest collected no tests from ${selection}; passing"
             exit 0
           fi
           exit "$status"
@@ -226,7 +251,7 @@ jobs:
   upload-coverage:
     name: Upload coverage to Codecov
     needs: run
-    if: always() && needs.run.outputs.decision != 'skip'
+    if: always() && needs.run.outputs.decision != 'skip' && needs.run.outputs.has-coverage == 'true'
     runs-on: ubuntu-latest
     permissions:
       contents: read

.github/workflows/test-unit.yml

@@ -35,6 +35,10 @@ concurrency:
 # already a matrix and carries a shard-coverage guard that reads that file by
 # name. Folding it in here is a follow-up, together with generalising that guard
 # into assert_ci_coverage.py.
+#
+# `fork-flag` names the `.circleci/tests.yml` job that now runs part of the
+# shard under the same Codecov flag. CircleCI does not build pull requests from
+# forks, so the shard still runs those files there and skips them elsewhere.
 jobs:
   unit:
     name: ${{ matrix.shard }}
@@ -65,10 +69,10 @@ jobs:
           - shard: enterprise-routing
             artifact-name: enterprise-routing
             test-path: >-
-              tests/test_litellm/enterprise
               tests/test_litellm/google_genai
               tests/test_litellm/router_utils
               tests/test_litellm/router_strategy
+            fork-flag: enterprise-routing
             workers: 2
             reruns: 2
             timeout-minutes: 20
@@ -200,7 +204,7 @@ jobs:
               tests/test_litellm/proxy/types_utils
               tests/test_litellm/proxy/logging_endpoints
               tests/test_litellm/proxy/test_*.py
-              tests/test_gateway
+            fork-flag: proxy-infra
             workers: 4
             reruns: 2
             timeout-minutes: 20
@@ -208,11 +212,8 @@ jobs:
 
           - shard: caching-local
             artifact-name: caching-local
-            test-path: >-
-              tests/local_testing/test_cache_preset_key.py
-              tests/local_testing/test_caching_handler.py
-              tests/local_testing/test_responses_stream_cache_keys.py
-              tests/local_testing/test_unit_test_caching.py
+            test-path: ""
+            fork-flag: caching-local
             workers: 2
             reruns: 2
             timeout-minutes: 20
@@ -220,7 +221,8 @@ jobs:
 
           - shard: proxy-extras
             artifact-name: proxy-extras
-            test-path: "tests/litellm-proxy-extras"
+            test-path: ""
+            fork-flag: proxy-extras
             workers: 2
             reruns: 2
             timeout-minutes: 20
@@ -228,7 +230,8 @@ jobs:
 
           - shard: enterprise-package
             artifact-name: enterprise-package
-            test-path: "tests/enterprise"
+            test-path: ""
+            fork-flag: enterprise-package
             workers: 4
             reruns: 2
             timeout-minutes: 20
@@ -247,6 +250,7 @@ jobs:
     uses: ./.github/workflows/_test-unit-base.yml
     with:
       test-path: ${{ matrix.test-path }}
+      fork-flag: ${{ matrix.fork-flag || '' }}
       workers: ${{ matrix.workers }}
       reruns: ${{ matrix.reruns }}
       timeout-minutes: ${{ matrix.timeout-minutes }}

Makefile

@@ -332,7 +332,7 @@ test-unit-core-utils: install-test-deps
 	$(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20
 
 test-unit-other: install-test-deps
-	$(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20
+	$(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20
 
 test-unit-root: install-test-deps
 	$(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20

tests/test_litellm/test_assert_ci_coverage.py

@@ -340,7 +340,7 @@ def test_a_dockerfile_directory_entry_is_stale_because_only_an_exact_path_exempt
 def test_a_workflow_that_names_a_file_clears_it_from_the_slice_check():
     named = coverage._workflow_named_tokens()
     assert named, "the workflows must name some test paths or the check proves nothing"
-    assert any(coverage._token_covers(token, "tests/local_testing/test_caching_handler.py") for token in named)
+    assert any(coverage._token_covers(token, "tests/proxy_unit_tests/test_proxy_custom_logger.py") for token in named)
 
 
 def test_the_slice_check_credits_only_workflows_never_the_circleci_config():

tests/unit/conftest.py

@@ -11,7 +11,7 @@ import litellm  # noqa: E402  # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im
 import litellm.router as litellm_router_module  # noqa: E402  # same import-time dependency
 import litellm.utils as litellm_utils_module  # noqa: E402  # same import-time dependency
 
-LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"]
+LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"]
 AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = (
     "AZURE_AD_TOKEN",
     "AZURE_TENANT_ID",

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -9,7 +9,7 @@ import pytest
 sys.path.insert(
     0,
     os.path.abspath(
-        os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras")
+        os.path.join(os.path.dirname(__file__), "../../../litellm-proxy-extras")
     ),
 )
 
@@ -23,7 +23,7 @@ from litellm_proxy_extras.utils import (
 _MIGRATIONS_DIR = os.path.abspath(
     os.path.join(
         os.path.dirname(__file__),
-        "../../litellm-proxy-extras/litellm_proxy_extras/migrations",
+        "../../../litellm-proxy-extras/litellm_proxy_extras/migrations",
     )
 )
 
@@ -999,7 +999,7 @@ class TestJWTKeyMappingCascade:
         schema_paths = glob.glob(
             os.path.abspath(
                 os.path.join(
-                    os.path.dirname(__file__), "../../**/schema.prisma"
+                    os.path.dirname(__file__), "../../../**/schema.prisma"
                 )
             ),
             recursive=True,