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

* 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: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests

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

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

* build: point the local proxy unit targets at the nested tests/unit/proxy tree

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

---------

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

.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 @@ GLOB_CHARS = frozenset("*?")
 # 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"
     )

tests/test_litellm/test_circleci_path_filter.py

@@ -107,7 +107,7 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"]
         ),
         (
             "cost-map-only",
-            ["model_prices_and_context_window.json", "tests/proxy_unit_tests/test_y.py"],
+            ["model_prices_and_context_window.json", "tests/unit/proxy/test_y.py"],
             "run",
         ),
         (

tests/unit/proxy/conftest.py

@@ -0,0 +1,150 @@
+# conftest.py
+
+import asyncio
+import copy
+import inspect
+import warnings
+
+import pytest
+
+
+import litellm
+import litellm.proxy.proxy_server
+
+
+# Top-level assignments of these types are the ones importlib.reload(litellm)
+# would have effectively reset. We snapshot them at conftest import time and
+# deep-copy the snapshot back before every test.
+_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes)
+
+
+def _snapshot_mutable_state(module):
+    """Capture a per-module snapshot of primitive and collection attributes."""
+    snapshot = {}
+    for attr in list(vars(module)):
+        if attr.startswith("_"):
+            continue
+        try:
+            value = getattr(module, attr)
+        except Exception as exc:
+            warnings.warn(
+                f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}",
+                stacklevel=2,
+            )
+            continue
+        if value is None or isinstance(value, _SNAPSHOT_TYPES):
+            try:
+                snapshot[attr] = copy.deepcopy(value)
+            except Exception as exc:
+                warnings.warn(
+                    f"conftest: could not snapshot {module.__name__}.{attr}: {exc}",
+                    stacklevel=2,
+                )
+    return snapshot
+
+
+def _restore_mutable_state(module, snapshot):
+    for attr, default in snapshot.items():
+        try:
+            setattr(module, attr, copy.deepcopy(default))
+        except Exception as exc:
+            warnings.warn(
+                f"conftest: could not restore {module.__name__}.{attr}: {exc}",
+                stacklevel=2,
+            )
+
+
+def _collect_flushable_caches():
+    """Return (module, attr) pairs whose values expose flush_cache()."""
+    targets = []
+    for module in (litellm, litellm.proxy.proxy_server):
+        for attr in list(vars(module)):
+            if attr.startswith("_"):
+                continue
+            try:
+                value = getattr(module, attr)
+            except Exception:
+                continue
+            # Only instances — a class reference has an unbound flush_cache
+            # that can't be called without a self argument.
+            if inspect.isclass(value) or inspect.ismodule(value):
+                continue
+            if callable(getattr(value, "flush_cache", None)):
+                targets.append((module, attr))
+    return targets
+
+
+def _flush_caches(targets):
+    for module, attr in targets:
+        try:
+            value = getattr(module, attr)
+        except Exception:
+            continue
+        flush = getattr(value, "flush_cache", None)
+        if callable(flush):
+            try:
+                flush()
+            except Exception as exc:
+                warnings.warn(
+                    f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}",
+                    stacklevel=2,
+                )
+
+
+# Snapshot once at conftest import — these are the "clean" module states.
+_LITELLM_STATE = _snapshot_mutable_state(litellm)
+_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server)
+_FLUSHABLE_CACHES = _collect_flushable_caches()
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_and_teardown():
+    """Reset mutable module state on litellm and proxy_server before each test.
+
+    Replaces a previous importlib.reload(litellm) approach that cost ~17s
+    per test (re-executing the full litellm __init__ import chain).
+
+    What IS reset:
+      - Top-level module attributes of type list / dict / set / tuple
+        / str / int / float / bool / bytes, and None-valued attributes.
+        These cover callback lists, general_settings, master_key,
+        premium_user, prisma_client, etc. — anything the old reload() reset
+        by re-executing the module body.
+      - Any module-level object instance that exposes flush_cache() (the
+        DualCache and LLMClientCache family), which handles cache state
+        that can't round-trip through deepcopy because of internal locks.
+
+    What is NOT reset:
+      - Class instances without flush_cache() (e.g. ProxyLogging,
+        JWTHandler, FastAPI routers, loggers). If a test mutates such an
+        instance in-place (setattr on the instance, appending to one of
+        its internal lists, etc.), the mutation will leak into later tests.
+        Use pytest's monkeypatch.setattr() or a local fixture for those
+        cases — don't rely on this autouse fixture to undo them.
+    """
+    _restore_mutable_state(litellm, _LITELLM_STATE)
+    _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE)
+    _flush_caches(_FLUSHABLE_CACHES)
+
+    loop = asyncio.get_event_loop_policy().new_event_loop()
+    asyncio.set_event_loop(loop)
+    try:
+        yield
+    finally:
+        loop.close()
+        asyncio.set_event_loop(None)
+
+
+def pytest_collection_modifyitems(config, items):
+    # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests
+    custom_logger_tests = [
+        item for item in items if "custom_logger" in item.parent.name
+    ]
+    other_tests = [item for item in items if "custom_logger" not in item.parent.name]
+
+    # Sort tests based on their names
+    custom_logger_tests.sort(key=lambda x: x.name)
+    other_tests.sort(key=lambda x: x.name)
+
+    # Reorder the items list
+    items[:] = custom_logger_tests + other_tests

tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py

@@ -5,7 +5,7 @@ from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
 
 # Add project root to sys.path
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))
 
 from litellm.proxy.auth.user_api_key_auth import (
     _resolve_jwt_to_virtual_key,

tests/unit/proxy/test_proxy_config_unit_test.py

@@ -31,7 +31,7 @@ async def test_basic_reading_configs_from_files():
     example_config_yaml_path = os.path.join(current_path, "example_config_yaml")
 
     # get all the files from example_config_yaml
-    files = os.listdir(example_config_yaml_path)
+    files = [f for f in os.listdir(example_config_yaml_path) if f.endswith((".yaml", ".yml"))]
     print(files)
 
     for file in files:

tests/unit/proxy/test_proxy_server.py

@@ -477,7 +477,7 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth):
     assert e.code == str(403)
 
 
-from test_custom_callback_input import CompletionCustomHandler
+from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler
 
 
 @mock_patch_acompletion()
@@ -1114,7 +1114,7 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
 from litellm.proxy.management_endpoints.team_endpoints import team_member_add
-from test_key_generate_prisma import prisma_client
+from tests.unit.proxy.management_endpoints.test_key_generate_prisma import prisma_client
 
 
 @pytest.fixture

tests/unit/skills/test_skills_db.py

@@ -42,7 +42,7 @@ def create_skill_zip(skill_name: str):
 
     The zip file is automatically cleaned up after use.
     """
-    test_dir = Path(__file__).parent.parent / "llm_translation" / "test_skills_data"
+    test_dir = Path(__file__).parents[2] / "llm_translation" / "test_skills_data"
     skill_dir = test_dir / skill_name
 
     # Create a zip file containing the skill directory

tests/unit/skills/test_skills_main.py

@@ -30,7 +30,7 @@ def test_create_skill_forwards_description_and_instructions_from_top_level_kwarg
 
 
 def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None:
-    """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under
+    """The SDK convention (see tests/unit/skills/test_skills_db.py) nests them under
     extra_body instead of passing them as top-level kwargs; both paths must reach the DB."""
     handler = MagicMock()
     monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler)