diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 87e7e17feb7..a631a7c3005 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -31,8 +31,15 @@ jobs: test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py" workers: 8 timeout: 20 + # test_proxy_utils.py is large (168+ parametrized tests) — run it on its + # own matrix so --dist=loadscope doesn't pin all of it to a single xdist + # worker and push the "remaining" group past the job timeout. + - test-group: proxy-utils + test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + workers: 8 + timeout: 20 - test-group: remaining - test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py" + test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --ignore=tests/proxy_unit_tests/test_proxy_utils.py" workers: 8 timeout: 30 uses: ./.github/workflows/_test-unit-services-base.yml diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 1421700c9a8..a0326f64ed7 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -1,48 +1,143 @@ # conftest.py -import importlib +import asyncio +import copy +import inspect import os import sys +import warnings import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path + 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. """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. - """ - curr_dir = os.getcwd() # Get the current working directory - sys.path.insert( - 0, os.path.abspath("../..") - ) # Adds the project directory to the system path - - import litellm - from litellm import Router - - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - import asyncio + _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) - print(litellm) - # from litellm import Router, completion, aembedding, acompletion, embedding - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) def pytest_collection_modifyitems(config, items):