ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests

.circleci/scripts/classify_changes.sh

@@ -31,7 +31,7 @@ while IFS= read -r file || [ -n "$file" ]; do
   case "$file" in
     model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
       has_cost_map=true ;;
-    tests/test_litellm/* | tests/proxy_unit_tests/*) : ;;
+    tests/test_litellm/* | tests/proxy_unit_tests/* | tests/unit/proxy/*) : ;;
     *) outside_cost_map_set=true ;;
   esac
 done

.circleci/scripts/unit_selection.sh

@@ -7,6 +7,18 @@ legacy_flags=(
   caching-local
   enterprise-package
   enterprise-routing
+  proxy-db-auth-checks
+  proxy-db-budgets
+  proxy-db-custom-logging
+  proxy-db-db-and-spend
+  proxy-db-endpoints-and-responses
+  proxy-db-guardrails-hooks
+  proxy-db-jwt-and-keys
+  proxy-db-key-generation
+  proxy-db-logging-misc
+  proxy-db-proxy-runtime
+  proxy-db-proxy-server-core
+  proxy-db-proxy-utils
   proxy-extras
   proxy-infra
 )
@@ -34,6 +46,66 @@ legacy_paths() {
       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-db-auth-checks)
+      echo tests/unit/proxy/auth/test_auth_checks.py
+      echo tests/unit/proxy/auth/test_user_api_key_auth.py
+      echo tests/unit/proxy/test_deprecated_key_grace_period.py ;;
+    proxy-db-budgets)
+      echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py
+      echo tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
+      echo tests/unit/proxy/test_zero_cost_model_budget_bypass.py ;;
+    proxy-db-custom-logging)
+      echo tests/unit/proxy/test_custom_callback_input.py
+      echo tests/unit/proxy/test_custom_logger_s3_gcs.py ;;
+    proxy-db-db-and-spend)
+      echo tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py
+      echo tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py
+      echo tests/unit/proxy/db/test_update_daily_tag_spend.py
+      echo tests/unit/proxy/test_db_schema_changes.py
+      echo tests/unit/proxy/test_prisma_client_backoff_retry.py
+      echo tests/unit/proxy/test_update_spend.py
+      echo tests/unit/skills/test_skills_db.py ;;
+    proxy-db-endpoints-and-responses)
+      echo tests/unit/proxy/auth/test_models_fallback_endpoint.py
+      echo tests/unit/proxy/common_utils/test_check_batch_cost.py
+      echo tests/unit/proxy/common_utils/test_check_responses_cost.py
+      echo tests/unit/proxy/common_utils/test_realtime_cache.py
+      echo tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py
+      echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
+      echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
+      echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
+      echo tests/unit/proxy/response_polling/test_response_polling_handler.py
+      echo tests/unit/proxy/test_custom_tokenizer_bug.py
+      echo tests/unit/proxy/test_get_favicon.py
+      echo tests/unit/proxy/test_get_image.py
+      echo tests/unit/proxy/test_prompt_test_endpoint.py
+      echo tests/unit/proxy/test_reducto_ocr_route.py
+      echo tests/unit/proxy/test_response_polling_pre_call_checks.py
+      echo tests/unit/proxy/test_ui_path_detection.py ;;
+    proxy-db-guardrails-hooks)
+      echo tests/unit/proxy/hooks/test_banned_keyword_list.py
+      echo tests/unit/proxy/test_proxy_setting_guardrails.py
+      echo tests/unit/proxy/test_unit_test_proxy_hooks.py ;;
+    proxy-db-jwt-and-keys)
+      echo tests/unit/proxy/auth/test_jwt.py
+      echo tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py
+      echo tests/unit/proxy/test_proxy_custom_auth.py ;;
+    proxy-db-key-generation) echo tests/unit/proxy/management_endpoints/test_key_generate_prisma.py ;;
+    proxy-db-logging-misc)
+      echo tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
+      echo tests/unit/proxy/spend_tracking/test_search_api_logging.py
+      echo tests/unit/proxy/test_proxy_reject_logging.py ;;
+    proxy-db-proxy-runtime)
+      echo tests/unit/proxy/auth/test_multipart_bypass_repro.py
+      echo tests/unit/proxy/auth/test_proxy_routes.py
+      echo tests/unit/proxy/middleware/test_request_size_limit_middleware.py
+      echo tests/unit/proxy/test_proxy_config_unit_test.py
+      echo tests/unit/proxy/test_proxy_token_counter.py
+      echo tests/unit/proxy/test_server_root_path.py ;;
+    proxy-db-proxy-server-core)
+      echo tests/unit/proxy/test_aproxy_startup.py
+      echo tests/unit/proxy/test_proxy_server.py ;;
+    proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.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 ;;

.circleci/tests.yml

@@ -317,7 +317,35 @@ workflows:
           reruns: 2
           matrix:
             parameters:
-              flag: [enterprise-package, proxy-infra]
+              flag:
+                - enterprise-package
+                - proxy-infra
+                - proxy-db-auth-checks
+                - proxy-db-jwt-and-keys
+                - proxy-db-proxy-server-core
+                - proxy-db-proxy-runtime
+                - proxy-db-custom-logging
+                - proxy-db-logging-misc
+                - proxy-db-db-and-spend
+                - proxy-db-guardrails-hooks
+                - proxy-db-budgets
+                - proxy-db-endpoints-and-responses
+          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-proxy-db-proxy-utils
+          flag: proxy-db-proxy-utils
+          shards: 1
+          reruns: 2
+          dist: worksteal
+          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-proxy-db-key-generation
+          flag: proxy-db-key-generation
+          shards: 1
+          workers: 0
+          reruns: 2
           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

.github/scripts/assert_ci_coverage.py

@@ -34,7 +34,6 @@
 # tests has to be named by some shard or it runs nowhere. A child listed here is
 # itself decomposed one level deeper and is checked through its own entry.
 SHARDED_ROOTS: tuple[str, ...] = (
-    "tests/proxy_unit_tests",
     "tests/test_litellm",
     "tests/test_litellm/proxy",
 )

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

@@ -20,6 +20,12 @@ concurrency:
 # 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. CircleCI does not build pull requests from forks, so `fork-flag`
+# makes the shard run that list there. `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 +
@@ -58,7 +64,7 @@ jobs:
   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)"
+    # 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:
@@ -71,132 +77,93 @@ jobs:
         include:
           # 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"
+            test-path: ""
+            fork-flag: proxy-db-key-generation
             workers: 0
             dist: loadscope
             timeout: 20
 
           # ---- 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
-              tests/proxy_unit_tests/test_deprecated_key_grace_period.py
+            test-path: ""
+            fork-flag: proxy-db-auth-checks
             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
+            test-path: ""
+            fork-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: "tests/proxy_unit_tests/test_proxy_utils.py"
+            test-path: ""
+            fork-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.py
-              tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py
-              tests/proxy_unit_tests/test_aproxy_startup.py
+            test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py"
+            fork-flag: proxy-db-proxy-server-core
             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_server_root_path.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
+            test-path: ""
+            fork-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_custom_callback_input.py
-              tests/proxy_unit_tests/test_custom_logger_s3_gcs.py
-              tests/proxy_unit_tests/test_proxy_custom_logger.py
+            test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py"
+            fork-flag: proxy-db-custom-logging
             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
+            test-path: ""
+            fork-flag: proxy-db-logging-misc
             workers: 4
             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_proxy_encrypt_decrypt.py
+            test-path: ""
+            fork-flag: proxy-db-db-and-spend
             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
+            test-path: ""
+            fork-flag: proxy-db-guardrails-hooks
             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
+            test-path: ""
+            fork-flag: proxy-db-budgets
             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_gemini_agents_endpoints.py
-              tests/proxy_unit_tests/test_get_favicon.py
-              tests/proxy_unit_tests/test_get_image.py
-              tests/proxy_unit_tests/test_reducto_ocr_route.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
+            test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py"
+            fork-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 }}
+      fork-flag: ${{ matrix.fork-flag }}
       workers: ${{ matrix.workers }}
       reruns: 2
       timeout-minutes: ${{ matrix.timeout }}

.github/workflows/test-unit.yml

@@ -31,10 +31,10 @@ concurrency:
 # number, so a partially-specified entry would fail the call rather than fall
 # back to the default.
 #
-# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is
-# 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.
+# tests/unit/proxy keeps its own caller (test-unit-proxy-db.yml): it is 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

Makefile

@@ -51,8 +51,8 @@ help:
 	@echo "  make test-unit-core-utils - Run core utils tests (~32 files)"
 	@echo "  make test-unit-other    - Run other tests (caching, responses, etc., ~69 files)"
 	@echo "  make test-unit-root     - Run root-level tests (~34 files)"
-	@echo "  make test-proxy-unit-a  - Run proxy_unit_tests (a-o, ~20 files)"
-	@echo "  make test-proxy-unit-b  - Run proxy_unit_tests (p-z, ~28 files)"
+	@echo "  make test-proxy-unit-a  - Run tests/unit/proxy (a-o)"
+	@echo "  make test-proxy-unit-b  - Run tests/unit/proxy (p-z)"
 	@echo "  make test-integration   - Run integration tests"
 	@echo "  make test-unit-helm     - Run helm unit tests"
 	@echo "  make test-rust-extension - Build the Rust extension and run its public Python tests"
@@ -337,12 +337,12 @@ test-unit-other: install-test-deps
 test-unit-root: install-test-deps
 	$(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
 
-# Proxy unit tests (tests/proxy_unit_tests split alphabetically)
+# Proxy unit tests (tests/unit/proxy split alphabetically)
 test-proxy-unit-a: install-test-deps
-	$(UV_RUN) pytest tests/proxy_unit_tests/test_[a-o]*.py --tb=short -vv -n 2 --durations=20
+	$(UV_RUN) pytest tests/unit/proxy --ignore-glob='tests/unit/proxy/test_[p-z]*.py' --tb=short -vv -n 2 --durations=20
 
 test-proxy-unit-b: install-test-deps
-	$(UV_RUN) pytest tests/proxy_unit_tests/test_[p-z]*.py --tb=short -vv -n 2 --durations=20
+	$(UV_RUN) pytest tests/unit/proxy/test_[p-z]*.py tests/unit/skills --tb=short -vv -n 2 --durations=20
 
 test-integration: install-test-deps
 	$(UV_RUN) pytest tests/ -k "not test_litellm"

litellm/llms/litellm_proxy/skills/README.md

@@ -369,7 +369,7 @@ model LiteLLM_SkillsTable {
 Run the tests:
 
 ```bash
-pytest tests/proxy_unit_tests/test_skills_db.py -v
+pytest tests/unit/skills/test_skills_db.py -v
 ```
 
 Tests cover:

tests/code_coverage_tests/user_api_key_auth_code_coverage.py

@@ -31,11 +31,11 @@ def get_function_names_from_file(file_path):
 def get_all_functions_called_in_tests(base_dir):
     """
     Returns a set of function names that are called in test functions
-    inside 'local_testing' and 'proxy_unit_tests' directories,
+    inside 'local_testing' and 'unit/proxy' directories,
     specifically in files containing the word 'router'.
     """
     called_functions = set()
-    test_dirs = ["local_testing", "proxy_unit_tests"]
+    test_dirs = ["local_testing", "unit/proxy"]
 
     for test_dir in test_dirs:
         dir_path = os.path.join(base_dir, test_dir)

tests/test_litellm/proxy/image_endpoints/test_azure_routes.py

@@ -53,7 +53,8 @@ def client_no_auth():
     config_fp = (
         repo_root
         / "tests"
-        / "proxy_unit_tests"
+        / "unit"
+        / "proxy"
         / "test_configs"
         / "test_config_no_auth.yaml"
     )