litellm/tests/local_testing/conftest.py
yuneng-jiang 1a00dd4dbb Fix router test isolation for xdist and rebalance proxy unit tests
Router tests: expand conftest save/restore to cover all globals mutated
by router tests (default_fallbacks, tag_budget_config, request_timeout,
enable_azure_ad_token_refresh, num_retries_per_request, model_cost,
token_counter). These were leaking across xdist workers.

Proxy tests: move test_proxy_utils.py (169 parametrized) and
test_proxy_server.py (72 parametrized) from part2 to part1, balancing
~370 vs ~360 tests (was ~129 vs ~600).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-15 21:36:56 -07:00

137 lines
4.0 KiB
Python

# conftest.py
#
# xdist-compatible test isolation for local_testing tests.
# Pattern matches tests/test_litellm/conftest.py:
# - Function-scoped fixture saves/restores litellm globals (no reload)
# - Module-scoped fixture reloads only in single-process mode
import importlib
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
@pytest.fixture(scope="function", autouse=True)
def isolate_litellm_state():
"""
Per-function isolation fixture.
Saves and restores litellm callback/global state so tests don't leak
side effects. Works safely under pytest-xdist parallel execution.
"""
# Save original callback state
original_state = {}
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
# Save other globals that tests commonly mutate
for attr in (
"set_verbose",
"cache",
"num_retries",
"num_retries_per_request",
"request_timeout",
"default_fallbacks",
"enable_azure_ad_token_refresh",
"tag_budget_config",
"model_cost",
"token_counter",
):
if hasattr(litellm, attr):
original_state[attr] = getattr(litellm, attr)
# Save rules that tests may set (e.g. test_rules.py)
for attr in ("pre_call_rules", "post_call_rules"):
if hasattr(litellm, attr):
val = getattr(litellm, attr)
original_state[attr] = val.copy() if val else []
# Save transport/network globals
for attr in ("disable_aiohttp_transport", "force_ipv4"):
if hasattr(litellm, attr):
original_state[attr] = getattr(litellm, attr)
# Flush cache before test
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
# Clear callbacks and rules before test
for attr in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
"pre_call_rules",
"post_call_rules",
):
if hasattr(litellm, attr):
setattr(litellm, attr, [])
yield
# Restore all saved state
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
for attr, original_value in original_state.items():
if hasattr(litellm, attr):
setattr(litellm, attr, original_value)
@pytest.fixture(scope="module", autouse=True)
def setup_and_teardown():
"""
Module-scoped setup. Reloads litellm only in single-process mode
(skipped under xdist to avoid cross-worker interference).
"""
sys.path.insert(0, os.path.abspath("../.."))
import litellm
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
if worker_id is None:
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
import litellm.proxy.proxy_server
importlib.reload(litellm.proxy.proxy_server)
except Exception as e:
print(f"Error reloading litellm.proxy.proxy_server: {e}")
if hasattr(litellm, "in_memory_llm_clients_cache"):
litellm.in_memory_llm_clients_cache.flush_cache()
yield
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