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

.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

@@ -74,6 +74,7 @@ commands:
     steps:
       - run:
           name: Install Codecov CLI (pinned v11.3.1)
+          when: always
           command: |
             curl -sSLf -o /tmp/codecov https://cli.codecov.io/v11.3.1/linux/codecov
             curl -sSLf -o /tmp/codecov.SHA256SUM https://cli.codecov.io/v11.3.1/linux/codecov.SHA256SUM
@@ -90,7 +91,6 @@ commands:
             uv run --no-sync python -c "import litellm_enterprise; print('litellm-enterprise OK:', litellm_enterprise.__file__)"
   setup_test_deps:
     steps:
-      - checkout
       - install_uv
       - install_rust
       - restore_cache:
@@ -165,42 +165,59 @@ commands:
 jobs:
   unit:
     parameters:
-      tests_path:
-        type: string
-        default: tests/unit
       flag:
         type: string
         default: unit
       shards:
         type: integer
         default: 6
+      workers:
+        type: integer
+        default: 4
+      dist:
+        type: string
+        default: loadscope
       base_ref:
         type: string
         default: ""
       pull_request_url:
         type: string
         default: ""
+      reruns:
+        type: integer
+        default: 0
     machine:
       image: ubuntu-2204:2024.04.1
     resource_class: large
     working_directory: ~/project
     parallelism: << parameters.shards >>
     environment:
+      COVERAGE_CORE: sysmon
       LITELLM_LOCAL_MODEL_COST_MAP: "True"
     steps:
-      - setup_test_deps
+      - checkout
       - skip_unless_relevant:
           base_ref: << parameters.base_ref >>
           pull_request_url: << parameters.pull_request_url >>
+      - setup_test_deps
       - run:
-          name: "Run << parameters.tests_path >> shard"
+          name: "Run << parameters.flag >> shard"
           no_output_timeout: 20m
           command: |
             mkdir -p test-results/<< parameters.flag >>
-            mapfile -t files < <(find << parameters.tests_path >> -name 'test_*.py' | sort | circleci tests split --split-by=timings --timings-type=filename)
-            if [ "${#files[@]}" -eq 0 ]; then echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.tests_path >> files; nothing to run"; exit 0; fi
+            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; }
+            mapfile -t files < <(printf '%s\n' "${shard}")
+            if [ "${#files[@]}" -eq 0 ]; then echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; fi
+            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
-            uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -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
+            env -i "${test_env[@]}" \
+              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
@@ -224,6 +241,7 @@ jobs:
     resource_class: large
     working_directory: ~/project
     steps:
+      - checkout
       - setup_test_deps
       - run:
           name: Checkout litellm-docs
@@ -250,16 +268,17 @@ jobs:
     resource_class: large
     working_directory: ~/project
     steps:
-      - setup_test_deps
+      - checkout
       - skip_unless_relevant:
           base_ref: << parameters.base_ref >>
           pull_request_url: << parameters.pull_request_url >>
+      - setup_test_deps
       - start_postgres:
           image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
       - start_redis
       - run:
           name: Run owned integration contracts
-          command: bash .circleci/scripts/run_integration.sh << parameters.suite >>
+          command: env -i PATH="$PATH" HOME="$HOME" CIRCLE_SHA1="$CIRCLE_SHA1" CIRCLE_WORKFLOW_ID="$CIRCLE_WORKFLOW_ID" bash .circleci/scripts/run_integration.sh << parameters.suite >>
           no_output_timeout: 15m
       - run:
           name: Stop owned database and Redis
@@ -282,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,20 +170,32 @@ 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 }}
           TEST_TIMEOUT_SECONDS: ${{ inputs.test-timeout-seconds }}
           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-proxy-db.yml

@@ -106,6 +106,7 @@ jobs:
           - test-group: proxy-server-core
             test-path: >-
               tests/proxy_unit_tests/test_proxy_server.py
+              tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py
               tests/proxy_unit_tests/test_aproxy_startup.py
             workers: 4
             dist: loadscope
@@ -115,7 +116,6 @@ jobs:
               tests/proxy_unit_tests/test_proxy_config_unit_test.py
               tests/proxy_unit_tests/test_proxy_routes.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
               tests/proxy_unit_tests/test_request_size_limit_middleware.py
               tests/proxy_unit_tests/test_multipart_bypass_repro.py