diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f26d93bced8..48a22dcc8af 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -467,76 +467,3 @@ async def test_e2e_generate_cold_storage_object_key_not_configured(): assert result is None -@pytest.mark.asyncio -async def test_logging_opentelemetry_context_propagation(): - """ - Test that OpenTelemtry context propagation works with async completion. - """ - import asyncio - - from opentelemetry import trace - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( - InMemorySpanExporter, - ) - - import litellm - from litellm.integrations.custom_logger import CustomLogger - - provider = TracerProvider() - exporter = InMemorySpanExporter() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - trace.set_tracer_provider(provider) - tracer = trace.get_tracer(__name__) - - class MockOpenTelemetryLogger(CustomLogger): - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - span = tracer.start_span(start_time=start_time.timestamp() * 1e9, name="async_log_success_event") - span.end(end_time=end_time.timestamp() * 1e9) - - - mock_logging_obj = MockOpenTelemetryLogger() - - litellm.callbacks = [mock_logging_obj] - - with tracer.start_as_current_span("span_1") as span: - span_1_id = span.get_span_context().span_id - await litellm.acompletion( - max_tokens=100, - messages=[{"role": "user", "content": "Hey"}], - model="openai/codex-mini-latest", - mock_response="Hello, world!", - ) - - - with tracer.start_as_current_span("span_2") as span: - span_2_id = span.get_span_context().span_id - await litellm.acompletion( - max_tokens=100, - messages=[{"role": "user", "content": "Hey"}], - model="openai/codex-mini-latest", - mock_response="Hello, world!", - ) - - await asyncio.sleep(1) - spans = exporter.get_finished_spans() - assert len(spans) == 4 - assert span_1_id != span_2_id - sorted_spans = sorted(list(spans), key=lambda x: x.start_time or 0) - - assert sorted_spans[0].name == "span_1" - assert sorted_spans[1].name == "async_log_success_event" - assert sorted_spans[2].name == "span_2" - assert sorted_spans[3].name == "async_log_success_event" - - first_span_context = sorted_spans[0].get_span_context() - assert first_span_context is not None and first_span_context.span_id == span_1_id - second_span_context = sorted_spans[2].get_span_context() - assert second_span_context is not None and second_span_context.span_id == span_2_id - first_completion_span_parent = sorted_spans[1].parent - assert first_completion_span_parent is not None and first_completion_span_parent.span_id == span_1_id - - # This check would fail without the proper context propagation, and span[3] would end up with span_1_id as the parent - second_completion_span_parent = sorted_spans[3].parent - assert second_completion_span_parent is not None and second_completion_span_parent.span_id == span_2_id