test(otel): wait for each exported span before the next request in Arize OTel v2 audit tests

This commit is contained in:
Devin AI 2026-09-30 03:46:59 +00:00
parent 4b0d28f704
commit e6f66c97b1
3 changed files with 76 additions and 34 deletions

View file

@ -561,55 +561,65 @@ def _matching_marker_spans(requests: tuple[Request, ...], marker: str) -> Iterat
yield attributes
def _single_span(spans: tuple[dict[str, str], ...]) -> dict[str, str]:
assert len(spans) == 1, spans
return spans[0]
def _matching_span(destination: Wire, response_id: str) -> dict[str, str]:
return eventually(
spans: Final = eventually(
lambda: tuple(_matching_llm_spans(destination.drain(), response_id)),
lambda spans: len(spans) == 1,
bool,
seconds=30,
)[0]
)
return _single_span(spans)
def _matching_marker_span(destination: Wire, marker: str) -> dict[str, str]:
return eventually(
spans: Final = eventually(
lambda: tuple(_matching_marker_spans(destination.drain(), marker)),
lambda spans: len(spans) == 1,
bool,
seconds=30,
)[0]
)
return _single_span(spans)
def _matching_output_value_span(destination: Wire, marker: str) -> dict[str, str]:
return eventually(
spans: Final = eventually(
lambda: tuple(
attributes
for attributes in _spans(destination.drain())
if attributes.get("openinference.span.kind") == "LLM" and marker in attributes.get("output.value", "")
),
lambda spans: len(spans) == 1,
bool,
seconds=30,
)[0]
)
return _single_span(spans)
def _matching_genai_marker_span(destination: Wire, marker: str) -> dict[str, str]:
return eventually(
spans: Final = eventually(
lambda: tuple(
attributes
for attributes in _spans(destination.drain())
if attributes.get("gen_ai.operation.name") == "chat" and marker in attributes.values()
),
lambda spans: len(spans) == 1,
bool,
seconds=30,
)[0]
)
return _single_span(spans)
def _matching_any_marker_span(destination: Wire, marker: str) -> dict[str, str]:
def matches(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]:
return tuple(attributes for attributes in _spans(requests) if marker in attributes.values())
return eventually(
spans: Final = eventually(
lambda: matches(destination.drain()),
lambda spans: len(spans) == 1,
bool,
seconds=30,
)[0]
)
return _single_span(spans)
def _collect_marker_spans(
@ -636,6 +646,20 @@ def _collect_marker_spans(
return collect(())
def _llm_spans_through_marker(destination: Wire, marker: str) -> tuple[dict[str, str], ...]:
def llm_spans(requests: tuple[Request, ...]) -> tuple[dict[str, str], ...]:
return tuple(
attributes for attributes in _spans(requests) if attributes.get("openinference.span.kind") == "LLM"
)
def collect(previous: tuple[dict[str, str], ...]) -> tuple[dict[str, str], ...]:
current: Final = eventually(lambda: llm_spans(destination.drain()), bool, seconds=30)
combined: Final = (*previous, *current)
return combined if any(marker in attributes.values() for attributes in combined) else collect(combined)
return collect(())
def _write_config(
directory: Path,
*,

View file

@ -27,6 +27,7 @@ from _openinference_support import (
_chat_tool_call,
_collect_marker_spans,
_json_object,
_matching_marker_span,
_messages_caller_raw_stream,
_messages_caller_response,
_normalize_chat_caller_stream,
@ -297,14 +298,7 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path
timer.cancel()
if timer.ident is not None:
timer.join(timeout=5)
requests: Final = rig.destination.drain()
spans: Final = tuple(
attributes
for attributes in _spans(requests)
if attributes.get("openinference.span.kind") == "LLM"
and attributes.get("litellm.metadata.trace_marker") == marker
)
assert len(spans) == 1, spans
_matching_marker_span(rig.destination, marker)
def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None:

View file

@ -24,11 +24,13 @@ from _openinference_support import (
_chat_caller_response,
_chat_caller_stream,
_chat_plain_response,
_chat_request_marker,
_chat_response,
_chat_stream_response,
_chat_tool_call,
_json_messages,
_json_object,
_llm_spans_through_marker,
_matching_marker_span,
_messages_caller_response,
_messages_caller_stream,
@ -43,7 +45,7 @@ from _openinference_support import (
_responses_stream_response,
_rig,
)
from integration._support.client import Gateway
from integration._support.client import Gateway, object_value, string_value
from integration._support.wire import Reply, Request
from pydantic import JsonValue
@ -844,27 +846,34 @@ def _cache_call(
@pytest.mark.parametrize("surface", ("chat", "chat-stream", "responses", "messages"))
def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -> None:
marker: Final = f"a-cache-{surface}-" + uuid.uuid4().hex
sentinel: Final = f"{marker}-sentinel"
def upstream(request: Request) -> Reply:
body: Final = _json_object(request.body)
request_marker: Final = (
_chat_request_marker(request)
if surface in ("chat", "chat-stream", "messages")
else string_value(object_value(body["metadata"])["trace_marker"])
)
assert request_marker in (marker, sentinel), request
if surface in ("chat", "chat-stream"):
_assert_chat_request(
request,
messages=[{"role": "user", "content": marker}],
messages=[{"role": "user", "content": request_marker}],
stream=True if surface == "chat-stream" else None,
stream_options={"include_usage": False} if surface == "chat-stream" else None,
)
return (
_chat_stream_response(marker, (_chat_tool_call(marker),), include_usage=False)
_chat_stream_response(request_marker, (_chat_tool_call(request_marker),), include_usage=False)
if surface == "chat-stream"
else _chat_response(marker)
else _chat_response(request_marker)
)
if surface == "responses":
assert body.get("metadata") == {"trace_marker": marker}, body
_assert_responses_request(request, marker=marker, input_value=marker)
return _responses_response(marker)
_assert_messages_request(request, marker=marker, prompt=marker)
return _messages_response(marker)
assert body.get("metadata") == {"trace_marker": request_marker}, body
_assert_responses_request(request, marker=request_marker, input_value=request_marker)
return _responses_response(request_marker)
_assert_messages_request(request, marker=request_marker, prompt=request_marker)
return _messages_response(request_marker)
with _rig(
gateway,
@ -879,7 +888,8 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
assert first[0] == (f"call_{marker}", "lookup_weather", expected_arguments), first
assert first[1].startswith("resp_") if surface == "responses" else first[1] == marker, first
assert not first[2].get("x-litellm-cache-key"), first[2]
rig.destination.drain()
first_span: Final = _matching_marker_span(rig.destination, marker)
_assert_tool_span_for_marker(first_span, marker)
second: Final = _cache_call(rig, surface, marker, cache_hit=True)
assert second[0] == first[0], second
assert second[1].startswith("resp_") if surface == "responses" else second[1] == first[1], second
@ -891,7 +901,6 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
assert not second[2].get("x-litellm-cache-key"), second[2]
else:
assert second[2].get("x-litellm-cache-key"), second[2]
_assert_tool_span_for_marker(_matching_marker_span(rig.destination, marker), marker)
forwarded_body: Final = _json_object(forwarded[0].body)
if surface in ("chat", "chat-stream"):
assert "metadata" not in forwarded_body, forwarded[0]
@ -899,3 +908,18 @@ def test_arize_otel_v2_a_cache(surface: str, gateway: Gateway, tmp_path: Path) -
assert forwarded_body["metadata"] == {"trace_marker": marker}, forwarded[0]
else:
assert forwarded_body["metadata"] == {}, forwarded[0]
sentinel_response: Final = _cache_call(rig, surface, sentinel)
assert not sentinel_response[2].get("x-litellm-cache-key"), sentinel_response[2]
sentinel_forwarded: Final = tuple(
request
for request in rig.provider.drain()
if request.method == "POST" and sentinel.encode() in request.body
)
assert len(sentinel_forwarded) == 1, sentinel_forwarded
spans: Final = _llm_spans_through_marker(rig.destination, sentinel)
sentinel_spans: Final = tuple(
attributes for attributes in spans if attributes.get("litellm.metadata.trace_marker") == sentinel
)
assert len(sentinel_spans) == 1, spans
_assert_tool_span_for_marker(sentinel_spans[0], sentinel)
assert not any(attributes.get("litellm.metadata.trace_marker") == marker for attributes in spans), spans