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

name: "Unit Tests: Proxy DB Operations"

on:
  pull_request:
    branches:
      - main
      - "litellm_**"
  push:
    branches:
      - main

permissions:
  contents: read

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

# 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.
#
# `.circleci/tests.yml` runs each group's files on same-repo events under the
# `proxy-db-<group>` Codecov flag; `.circleci/scripts/unit_selection.sh` holds
# the file lists. That pipeline is manual-only while the tests migrate, so
# `unit-flag` makes the shard run that list on every event. `test-path` keeps
# the files that still reach real providers and never left
# tests/proxy_unit_tests.
#
# 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.
#   * `timeout` bounds the pytest step only. Checkout, dependency install, and
#     Prisma client generation draw on a separate allowance in the base
#     workflow, so slow setup shows up as a slow job rather than as a
#     cancelled shard whose tests were passing.
#   * 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).
jobs:
  # Fast guard — fails the workflow when a test directory or file inside a sharded
  # tree is claimed by no shard. The semantic-shard design has no catch-all bucket,
  # so an unassigned child runs nowhere; assert_ci_coverage.py holds the tree list
  # and reads the same test-path keys the coverage census does.
  assert-shard-coverage:
    runs-on: ubuntu-latest
    timeout-minutes: 2
    permissions:
      contents: read
    steps:
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          persist-credentials: false
      - name: Assert every test directory and file is claimed by a shard
        run: python3 .github/scripts/assert_ci_coverage.py --shards

  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/unit/proxy/…, 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
      pull-requests: write
    strategy:
      fail-fast: false
      matrix:
        include:
          # Must run serially — event-loop conflict with the logging worker.
          - test-group: key-generation
            test-path: ""
            unit-flag: proxy-db-key-generation
            workers: 0
            dist: loadscope
            timeout: 20

          # ---- auth: split into 2 shards ----
          - test-group: auth-checks
            test-path: ""
            unit-flag: proxy-db-auth-checks
            workers: 4
            dist: loadscope
            timeout: 15
          - test-group: jwt-and-keys
            test-path: ""
            unit-flag: proxy-db-jwt-and-keys
            workers: 4
            dist: loadscope
            timeout: 15

          # ---- test_proxy_utils.py, single shard, worksteal distribution ----
          - test-group: proxy-utils
            test-path: ""
            unit-flag: proxy-db-proxy-utils
            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_gemini_pass_through.py"
            unit-flag: proxy-db-proxy-server-core
            workers: 4
            dist: loadscope
            timeout: 15
          - test-group: proxy-runtime
            test-path: ""
            unit-flag: proxy-db-proxy-runtime
            workers: 4
            dist: loadscope
            timeout: 15

          # ---- logging: split into 2 shards ----
          - test-group: custom-logging
            test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py"
            unit-flag: proxy-db-custom-logging
            workers: 4
            dist: loadscope
            timeout: 15
          - test-group: logging-misc
            test-path: ""
            unit-flag: proxy-db-logging-misc
            workers: 4
            dist: loadscope
            timeout: 15

          - test-group: db-and-spend
            test-path: ""
            unit-flag: proxy-db-db-and-spend
            workers: 4
            dist: loadscope
            timeout: 15

          # ---- guardrails + budget + hooks: split into 2 ----
          - test-group: guardrails-hooks
            test-path: ""
            unit-flag: proxy-db-guardrails-hooks
            workers: 4
            dist: loadscope
            timeout: 15
          - test-group: budgets
            test-path: ""
            unit-flag: proxy-db-budgets
            workers: 4
            dist: loadscope
            timeout: 15

          - test-group: endpoints-and-responses
            test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py"
            unit-flag: proxy-db-endpoints-and-responses
            workers: 4
            dist: loadscope
            timeout: 15
    uses: ./.github/workflows/_test-unit-base.yml
    with:
      test-path: ${{ matrix.test-path }}
      unit-flag: ${{ matrix.unit-flag }}
      workers: ${{ matrix.workers }}
      reruns: 2
      timeout-minutes: ${{ matrix.timeout }}
      dist: ${{ matrix.dist }}
      artifact-name: proxy-db-${{ matrix.test-group }}