diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 7efcc76073f..85015cd2edf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -120,6 +120,7 @@ RESTORED_GLOBALS: Final = ( "anthropic_models", "token_counter", "initialized_langfuse_clients", + "turn_off_message_logging", ) MODULE_LEVEL_CLIENTS: Final = ("module_level_client", "module_level_aclient") SESSION_CLIENTS: Final = ("base_llm_aiohttp_handler", "httpx_client", "aclient", "client") diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 9ba09085ce8..6f80ff99b3e 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -7007,35 +7007,31 @@ def validate_redacted_message_span_attributes(span: ReadableSpan) -> None: ) +async def _first_finished_span(span_exporter: InMemorySpanExporter) -> None: + while not span_exporter.get_finished_spans(): + await asyncio.sleep(0.01) + + @pytest.mark.usefixtures("unset_global_tracer_provider") @pytest.mark.asyncio @pytest.mark.parametrize("streaming", [True, False]) @pytest.mark.parametrize("global_redact", [True, False]) -async def test_awesome_otel_with_message_logging_off(streaming, global_redact): +async def test_awesome_otel_with_message_logging_off(streaming, global_redact, monkeypatch: pytest.MonkeyPatch): """ No content should be logged when message logging is off tests when litellm.turn_off_message_logging is set to True tests when OpenTelemetry(message_logging=False) is set """ - litellm.set_verbose = True - - # Clear exporter at the start to ensure clean state - exporter.clear() - - litellm.callbacks = [OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter))] + span_exporter: Final = InMemorySpanExporter() + OpenTelemetry(config=OpenTelemetryConfig(exporter=span_exporter)) if global_redact is False: - otel_logger = OpenTelemetry( - message_logging=False, config=OpenTelemetryConfig(exporter="console") - ) + otel_logger = OpenTelemetry(message_logging=False, config=OpenTelemetryConfig(exporter="console")) else: - # use global redaction - litellm.turn_off_message_logging = True + monkeypatch.setattr(litellm, "turn_off_message_logging", True) otel_logger = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) litellm.callbacks = [otel_logger] - litellm.success_callback = [] - litellm.failure_callback = [] response = await litellm.acompletion( model="gpt-4.1-mini", @@ -7043,27 +7039,14 @@ async def test_awesome_otel_with_message_logging_off(streaming, global_redact): mock_response="hi", stream=streaming, ) - print("response", response) if streaming is True: - async for chunk in response: - print("chunk", chunk) + async for _ in response: + pass - await asyncio.sleep(1) - spans = exporter.get_finished_spans() - print("spans", spans) - assert len(spans) == 1 - - _span = spans[0] - print("span attributes", _span.attributes) - - validate_redacted_message_span_attributes(_span) - - # clear in memory exporter - exporter.clear() - - if global_redact is True: - litellm.turn_off_message_logging = False + await asyncio.wait_for(_first_finished_span(span_exporter), timeout=LOGGING_WORKER_DRAIN_TIMEOUT_SECONDS) + (span,) = span_exporter.get_finished_spans() + validate_redacted_message_span_attributes(span) @pytest.mark.usefixtures("unset_global_tracer_provider", "drained_logging_worker") diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index 63db0eb5bba..93be480500d 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -14,6 +14,7 @@ from litellm import utils as litellm_utils_module from litellm.caching.caching import DualCache from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from tests.unit import conftest as unit_harness CANARY_MODEL = "conftest-isolation-canary-model" @@ -105,6 +106,15 @@ def test_live_router_membership_was_rolled_back(): assert _CanaryRouterHolder.router not in litellm_router_module._live_routers +def test_turn_off_message_logging_is_scoped_to_this_test(): + litellm.turn_off_message_logging = True + assert should_redact_message_logging({}) is True + + +def test_turn_off_message_logging_was_rolled_back(): + assert litellm.turn_off_message_logging is False + + def test_aws_cache_reset_sees_new_handlers_and_replaced_caches(monkeypatch: pytest.MonkeyPatch) -> None: module: Final = unit_harness.importlib.import_module("litellm.files.main") handler: Final = BaseAWSLLM()