diff --git a/tests/otel_tests/test_otel.py b/tests/otel_tests/test_otel.py index a9ed4461a1..a0f58dd5b8 100644 --- a/tests/otel_tests/test_otel.py +++ b/tests/otel_tests/test_otel.py @@ -104,10 +104,24 @@ async def test_chat_completion_check_otel_spans(): print("otel_spans: ", otel_spans) all_otel_spans = otel_spans["otel_spans"] - most_recent_parent = str(otel_spans["most_recent_parent"]) - print("Most recent OTEL parent: ", most_recent_parent) - print("\n spans grouped by parent: ", otel_spans["spans_grouped_by_parent"]) - parent_trace_spans = otel_spans["spans_grouped_by_parent"][most_recent_parent] + spans_grouped_by_parent = otel_spans["spans_grouped_by_parent"] + print("\n spans grouped by parent: ", spans_grouped_by_parent) + + # The GET /otel-spans request itself produces auth spans that beat + # the chat-completion spans on start_time, so `most_recent_parent` + # points at the wrong trace. Pick the chat-completion trace by + # content: it's the one carrying the full set of expected markers. + chat_completion_markers = { + "postgres", + "redis", + "raw_gen_ai_request", + "batch_write_to_db", + } + parent_trace_spans = next( + spans + for spans in spans_grouped_by_parent.values() + if chat_completion_markers.issubset(spans) + ) print("Parent trace spans: ", parent_trace_spans)