Files
litellm/tests/llm_translation/conftest.py
T

112 lines
3.5 KiB
Python

# conftest.py
#
# xdist-compatible test isolation for llm_translation tests.
# Mirrors the pattern in tests/local_testing/conftest.py:
# - Function-scoped fixture resets litellm globals to true defaults
# - Module-scoped reload 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
import asyncio
# ---------------------------------------------------------------------------
# Capture TRUE defaults at conftest import time (before test modules pollute).
# ---------------------------------------------------------------------------
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"set_verbose": getattr(litellm, "set_verbose", False),
"cache": getattr(litellm, "cache", None),
"allowed_fails": getattr(litellm, "allowed_fails", 3),
"disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False),
"force_ipv4": getattr(litellm, "force_ipv4", False),
"drop_params": getattr(litellm, "drop_params", None),
"modify_params": getattr(litellm, "modify_params", False),
"api_base": getattr(litellm, "api_base", None),
"api_key": getattr(litellm, "api_key", None),
"cohere_key": getattr(litellm, "cohere_key", None),
}
@pytest.fixture(scope="session")
def event_loop():
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
yield loop
loop.close()
@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown(event_loop): # Add event_loop as a dependency
curr_dir = os.getcwd()
sys.path.insert(0, os.path.abspath("../.."))
import litellm
# ---- Save current state (for teardown restore) ----
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 []
for attr in _SCALAR_DEFAULTS:
if hasattr(litellm, attr):
original_state[attr] = getattr(litellm, attr)
# ---- Reset to true defaults before the test ----
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
importlib.reload(litellm)
# Set the event loop from the fixture
asyncio.set_event_loop(event_loop)
yield
# ---- Teardown ----
for attr, original_value in original_state.items():
if hasattr(litellm, attr):
setattr(litellm, attr, original_value)
# Clean up any pending tasks
pending = asyncio.all_tasks(event_loop)
for task in pending:
task.cancel()
# Run the event loop until all tasks are cancelled
if pending:
event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
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