mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-27 01:22:18 +00:00
test(e2e): accept the otel cost write as a linked root trace (#42931)
* test(e2e): accept the otel cost write as a linked root trace Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): window and fail-closed the linked otel trace read-back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): poll the otel read-back without recursion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
e11c3f5815
commit
3ebf6a1fd8
2 changed files with 222 additions and 130 deletions
|
|
@ -2,8 +2,9 @@
|
|||
|
||||
Covers logging.otel.success.exports_metric: a successful non-streaming call must
|
||||
land at the OTEL destination as ONE connected trace - a single root SERVER span
|
||||
with the auth phase, db lookups, and cost write under it, and the gen-AI CLIENT
|
||||
span parented into the same tree. The regression this pins: the proxy publishing
|
||||
with the auth phase and db lookups under it, the gen-AI CLIENT span parented
|
||||
into the same tree, and the cost write either under it or as the root of its
|
||||
own trace linked back to the request span. The regression this pins: the proxy publishing
|
||||
the global TracerProvider before callbacks init made server spans export through
|
||||
a different provider than the preset's gen-AI spans, so the destination received
|
||||
the gen-AI span alone, dangling (fixed in #30590; verified failing at its parent
|
||||
|
|
@ -28,7 +29,7 @@ from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, OTEL_EXPORTER_
|
|||
from lifecycle import ResourceManager
|
||||
from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body
|
||||
from models import LiteLLMParamsBody
|
||||
from otel_client import JaegerSpan, JaegerTrace, OtelReader
|
||||
from otel_client import CallTraces, JaegerSpan, JaegerTrace, OtelReader, root_span
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
|
@ -78,12 +79,13 @@ def _chain_reaches(span_id: str, root_id: str, trace: JaegerTrace) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _assert_complete_trace(
|
||||
hits: list[JaegerTrace], *, route: str, genai_span: str, require_cost_span: bool = True
|
||||
) -> None:
|
||||
"""The enforced behavior: the destination holds exactly one trace for the
|
||||
call, rooted at the SERVER span, with auth/db/cost children and the gen-AI
|
||||
span all connected into that one tree - no dangling parent references."""
|
||||
def _assert_complete_trace(traces: CallTraces, *, route: str, genai_span: str, require_cost_span: bool = True) -> None:
|
||||
"""The enforced behavior: the destination holds exactly one call-id-tagged
|
||||
trace for the call, rooted at the SERVER span, with auth/db children and
|
||||
the gen-AI span all connected into that one tree - no dangling parent
|
||||
references - and the cost write either in that trace or as the root of
|
||||
its own trace linked FOLLOWS_FROM to the request SERVER span."""
|
||||
hits = traces.hits
|
||||
assert hits, (
|
||||
"no trace for this call arrived at the destination within the deadline "
|
||||
"(nothing tagged with its call id was found)"
|
||||
|
|
@ -119,8 +121,20 @@ def _assert_complete_trace(
|
|||
assert any(name.startswith(DB_SPAN_PREFIX) for name in names), (
|
||||
f"no db ('{DB_SPAN_PREFIX}*') span in the trace; spans: {names}"
|
||||
)
|
||||
if require_cost_span:
|
||||
assert COST_SPAN in names, f"cost write span {COST_SPAN!r} missing; spans: {names}"
|
||||
if require_cost_span and COST_SPAN not in names:
|
||||
cost_traces = [t for t in traces.linked if (r := root_span(t)) is not None and r.operation_name == COST_SPAN]
|
||||
assert len(cost_traces) == 1, (
|
||||
f"cost write span {COST_SPAN!r} reached neither the request trace nor its own "
|
||||
f"trace linked to the request SERVER span; request spans: {names}; "
|
||||
f"linked traces: {[(t.trace_id, t.span_names()) for t in traces.linked]}"
|
||||
)
|
||||
cost_root = root_span(cost_traces[0])
|
||||
assert cost_root is not None, f"cost write trace has no single root; spans: {cost_traces[0].span_names()}"
|
||||
link = next(ref for ref in cost_root.references if ref.span_id == root.span_id)
|
||||
assert link.ref_type == "FOLLOWS_FROM" and link.trace_id == trace.trace_id, (
|
||||
f"the cost write trace's root must reference the request SERVER span FOLLOWS_FROM, "
|
||||
f"got refType={link.ref_type!r} traceID={link.trace_id!r} (request trace {trace.trace_id})"
|
||||
)
|
||||
|
||||
genai = next((span for span in trace.spans if span.operation_name == genai_span), None)
|
||||
assert genai is not None, f"gen-AI span {genai_span!r} missing; spans: {names}"
|
||||
|
|
@ -131,9 +145,15 @@ def _assert_complete_trace(
|
|||
)
|
||||
|
||||
|
||||
def _settled_names(*, route: str, genai_span: str, require_cost_span: bool = True) -> set[str]:
|
||||
names = {f"POST {route}", f"auth {route}", genai_span}
|
||||
return (names | {COST_SPAN}) if require_cost_span else names
|
||||
def _poll(
|
||||
otel_reader: OtelReader, *, call_id: str, route: str, genai_span: str, require_cost_span: bool = True
|
||||
) -> CallTraces:
|
||||
return otel_reader.poll_traces_for_call(
|
||||
call_id=call_id,
|
||||
settled_names={f"POST {route}", f"auth {route}", genai_span},
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
linked_names=frozenset({COST_SPAN}) if require_cost_span else frozenset(),
|
||||
)
|
||||
|
||||
|
||||
def _tag(span: JaegerSpan, key: str) -> str | int | float | bool | None:
|
||||
|
|
@ -174,7 +194,7 @@ def one_served_genai_span(trace: JaegerTrace, genai_span: str) -> JaegerSpan:
|
|||
return served[0]
|
||||
|
||||
|
||||
def _assert_real_ttft(hits: list[JaegerTrace], *, genai_span: str) -> None:
|
||||
def _assert_real_ttft(hits: tuple[JaegerTrace, ...], *, genai_span: str) -> None:
|
||||
"""The enforced behavior: the gen-AI span for the attempt that served the
|
||||
stream records a TTFT that is a real measurement - present, numeric,
|
||||
positive, and strictly less than that span's own total duration. A TTFT of
|
||||
|
|
@ -286,9 +306,10 @@ class TestOtelTraceCompleteness:
|
|||
/chat/completions request produces one complete OTEL trace.
|
||||
|
||||
The trace should have a single server root span for the incoming request, with
|
||||
the authentication, database, and cost-recording work beneath it. The span for
|
||||
the actual model call must also belong to that same trace, rather than being
|
||||
exported separately with a missing parent.
|
||||
the authentication and database work beneath it. The span for the actual model
|
||||
call must also belong to that same trace, rather than being exported separately
|
||||
with a missing parent, and the cost-recording work must land either in that
|
||||
trace or in its own trace linked to it.
|
||||
|
||||
This matters because a split trace is easy to miss: all of the spans may still
|
||||
arrive, but the model call appears without the surrounding request context.
|
||||
|
|
@ -308,12 +329,8 @@ class TestOtelTraceCompleteness:
|
|||
outcome = first_ok(client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16))
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}")
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}")
|
||||
_assert_complete_trace(traces, route=route, genai_span=f"chat {MODEL}")
|
||||
|
||||
@pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["chat_completions"])
|
||||
@pytest.mark.otel_tls
|
||||
|
|
@ -337,11 +354,7 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
|
||||
hits: Final = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
hits: Final = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}")
|
||||
_assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}")
|
||||
|
||||
@pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["messages"])
|
||||
|
|
@ -352,8 +365,10 @@ class TestOtelTraceCompleteness:
|
|||
produces exactly one complete OTEL trace.
|
||||
|
||||
The trace must have a single root span named "POST /v1/messages". The
|
||||
authentication, database, cost-writing, and model-call spans must all belong to
|
||||
authentication, database, and model-call spans must all belong to
|
||||
the same trace and have valid parent relationships leading back to that root.
|
||||
The cost-writing span must land in the request trace or in its own trace
|
||||
linked to it.
|
||||
|
||||
The model-call span is expected to be named "chat <model>". The test fails if
|
||||
the request is split across multiple traces, if any span references a missing
|
||||
|
|
@ -370,12 +385,8 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=f"chat {MODEL}"),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=f"chat {MODEL}")
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=f"chat {MODEL}")
|
||||
_assert_complete_trace(traces, route=route, genai_span=f"chat {MODEL}")
|
||||
|
||||
@pytest.mark.covers("logging.otel.success.exports_metric", exercised_on=["responses"])
|
||||
def test_responses_exports_complete_trace(
|
||||
|
|
@ -385,12 +396,14 @@ class TestOtelTraceCompleteness:
|
|||
produces exactly one complete OTEL trace.
|
||||
|
||||
The trace must have a single root span named "POST /v1/responses". The
|
||||
authentication, database, cost-writing, and model-call spans must all belong to
|
||||
the same trace and have valid parent relationships leading back to that root.
|
||||
authentication, database, and model-call spans must all belong to the same
|
||||
trace and have valid parent relationships leading back to that root. The cost
|
||||
write finishes after the response, so it lands as the root of its own trace
|
||||
linked FOLLOWS_FROM to the request SERVER span.
|
||||
|
||||
The model-call span is expected to be named "chat <model>". The test fails if
|
||||
the request is split across multiple traces, if any span references a missing
|
||||
parent, or if the model-call span cannot be connected back to the root."""
|
||||
The model-call span is expected to be named "chat <model>". The test fails on
|
||||
a split request trace, a dangling parent, a disconnected model-call span, or
|
||||
a cost write that is neither in the request trace nor linked to it."""
|
||||
route = "/v1/responses"
|
||||
_assert_otel_destination_configured(client)
|
||||
|
||||
|
|
@ -405,12 +418,8 @@ class TestOtelTraceCompleteness:
|
|||
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
|
||||
|
||||
genai_span = f"chat {CHEAP_OPENAI_MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.exports_metric", exercised_on=["chat_completions"])
|
||||
def test_chat_completions_stream_exports_complete_trace(
|
||||
|
|
@ -418,8 +427,9 @@ class TestOtelTraceCompleteness:
|
|||
) -> None:
|
||||
"""A successful streamed `/chat/completions` request should export one
|
||||
complete OTEL trace. The trace must contain a single root `SERVER`
|
||||
span, with the auth, database, cost, and gen-AI `CLIENT` spans all
|
||||
connected back to that root.
|
||||
span, with the auth, database, and gen-AI `CLIENT` spans all
|
||||
connected back to that root, and the cost write in that trace or in
|
||||
its own trace linked to it.
|
||||
|
||||
Streaming has an additional lifecycle risk because the gen-AI span is
|
||||
closed by the stream-consumption path after the final chunk has
|
||||
|
|
@ -451,14 +461,10 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span)
|
||||
|
||||
served = one_served_genai_span(hits[0], genai_span)
|
||||
served = one_served_genai_span(traces.hits[0], genai_span)
|
||||
assert _tag(served, "litellm.request.streaming") is True, (
|
||||
"the gen-AI span must record litellm.request.streaming=true; its absence means "
|
||||
"the stream flag was dropped before the model call"
|
||||
|
|
@ -470,8 +476,9 @@ class TestOtelTraceCompleteness:
|
|||
) -> None:
|
||||
"""A successful streamed `/v1/messages` request should export one
|
||||
complete OTEL trace. The trace must contain a single root `SERVER`
|
||||
span, with the auth, database, cost, and gen-AI `CLIENT` spans all
|
||||
connected back to that root.
|
||||
span, with the auth, database, and gen-AI `CLIENT` spans all
|
||||
connected back to that root, and the cost write in that trace or in
|
||||
its own trace linked to it.
|
||||
|
||||
This endpoint has the same streaming lifecycle risk as
|
||||
`/chat/completions`: the gen-AI span is closed by the
|
||||
|
|
@ -503,14 +510,10 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span)
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span)
|
||||
|
||||
served = one_served_genai_span(hits[0], genai_span)
|
||||
served = one_served_genai_span(traces.hits[0], genai_span)
|
||||
assert _tag(served, "litellm.request.streaming") is True, (
|
||||
"the gen-AI span must record litellm.request.streaming=true; its absence means "
|
||||
"the stream flag was dropped before the model call"
|
||||
|
|
@ -555,14 +558,12 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {CHEAP_OPENAI_MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
traces = _poll(
|
||||
otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
|
||||
one_served_genai_span(hits[0], genai_span)
|
||||
one_served_genai_span(traces.hits[0], genai_span)
|
||||
|
||||
spend_row = client.poll_proxy_spend_for_key(key)
|
||||
assert spend_row is not None and spend_row.spend is not None and spend_row.spend > 0, (
|
||||
|
|
@ -609,12 +610,8 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_real_ttft(hits, genai_span=genai_span)
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span)
|
||||
_assert_real_ttft(traces.hits, genai_span=genai_span)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.records_ttft", exercised_on=["messages"])
|
||||
def test_messages_stream_records_real_ttft(
|
||||
|
|
@ -651,12 +648,8 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
)
|
||||
_assert_real_ttft(hits, genai_span=genai_span)
|
||||
traces = _poll(otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span)
|
||||
_assert_real_ttft(traces.hits, genai_span=genai_span)
|
||||
|
||||
@pytest.mark.covers("logging.otel.stream.records_ttft", exercised_on=["responses"])
|
||||
def test_responses_stream_records_real_ttft(
|
||||
|
|
@ -693,12 +686,10 @@ class TestOtelTraceCompleteness:
|
|||
)
|
||||
|
||||
genai_span = f"chat {CHEAP_OPENAI_MODEL}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
traces = _poll(
|
||||
otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False
|
||||
)
|
||||
_assert_real_ttft(hits, genai_span=genai_span)
|
||||
_assert_real_ttft(traces.hits, genai_span=genai_span)
|
||||
|
||||
@pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["chat_completions"])
|
||||
def test_failed_chat_completions_error_span_attributes(
|
||||
|
|
@ -745,18 +736,16 @@ class TestOtelTraceCompleteness:
|
|||
assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id"
|
||||
|
||||
genai_span = f"chat {model_name}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
traces = _poll(
|
||||
otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
|
||||
root = next(span for span in hits[0].spans if not span.references)
|
||||
root = next(span for span in traces.hits[0].spans if not span.references)
|
||||
assert str(_tag(root, "http.status_code")) == "401", (
|
||||
f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}"
|
||||
)
|
||||
genai = next(span for span in hits[0].spans if span.operation_name == genai_span)
|
||||
genai = next(span for span in traces.hits[0].spans if span.operation_name == genai_span)
|
||||
_assert_error_span_contract(genai)
|
||||
|
||||
@pytest.mark.covers("logging.otel.failure.exports_metric", exercised_on=["messages"])
|
||||
|
|
@ -803,16 +792,14 @@ class TestOtelTraceCompleteness:
|
|||
assert outcome.call_id is not None, "failed responses must still carry x-litellm-call-id"
|
||||
|
||||
genai_span = f"chat {model_name}"
|
||||
hits = otel_reader.poll_traces_for_call(
|
||||
call_id=outcome.call_id,
|
||||
settled_names=_settled_names(route=route, genai_span=genai_span, require_cost_span=False),
|
||||
settled_prefixes={DB_SPAN_PREFIX},
|
||||
traces = _poll(
|
||||
otel_reader, call_id=outcome.call_id, route=route, genai_span=genai_span, require_cost_span=False
|
||||
)
|
||||
_assert_complete_trace(hits, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
_assert_complete_trace(traces, route=route, genai_span=genai_span, require_cost_span=False)
|
||||
|
||||
root = next(span for span in hits[0].spans if not span.references)
|
||||
root = next(span for span in traces.hits[0].spans if not span.references)
|
||||
assert str(_tag(root, "http.status_code")) == "401", (
|
||||
f"the SERVER span must record the 401 the client received, got {_tag(root, 'http.status_code')!r}"
|
||||
)
|
||||
genai = next(span for span in hits[0].spans if span.operation_name == genai_span)
|
||||
genai = next(span for span in traces.hits[0].spans if span.operation_name == genai_span)
|
||||
_assert_error_span_contract(genai)
|
||||
|
|
|
|||
|
|
@ -10,6 +10,11 @@ so the completeness assertions see the whole tree. A failed query is a hard
|
|||
failure, never an empty result - an unreachable destination must not read as
|
||||
"the trace never arrived".
|
||||
|
||||
Service spans that end after the response (the cost write is one) carry no
|
||||
call id and land as the root of their own trace with a link back to the
|
||||
request span, so they are fetched by operation name and matched by that link
|
||||
to the request root rather than by the tag query.
|
||||
|
||||
External reads go through ``e2e_http`` (the only module allowed to call
|
||||
``requests.*``).
|
||||
"""
|
||||
|
|
@ -18,7 +23,9 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
|
@ -84,18 +91,57 @@ class JaegerTracesPage(BaseModel):
|
|||
|
||||
class _TracesQuery(BaseModel):
|
||||
service: str
|
||||
tags: str
|
||||
tags: str | None = None
|
||||
operation: str | None = None
|
||||
limit: int = 20
|
||||
lookback: str = "1h"
|
||||
start: int | None = None
|
||||
end: int | None = None
|
||||
|
||||
|
||||
def _ticks() -> Iterator[None]:
|
||||
while True:
|
||||
yield None
|
||||
time.sleep(POLL_INTERVAL)
|
||||
|
||||
|
||||
def _settled(trace: JaegerTrace, names: set[str], prefixes: set[str]) -> bool:
|
||||
present = set(trace.span_names())
|
||||
return names.issubset(present) and all(
|
||||
any(name.startswith(prefix) for name in present) for prefix in prefixes
|
||||
return names.issubset(present) and all(any(name.startswith(prefix) for name in present) for prefix in prefixes)
|
||||
|
||||
|
||||
def root_span(trace: JaegerTrace) -> JaegerSpan | None:
|
||||
"""The single span whose references all point outside the trace (a span
|
||||
with no references qualifies). None when there is not exactly one."""
|
||||
in_trace = {span.span_id for span in trace.spans}
|
||||
roots = [span for span in trace.spans if all(ref.span_id not in in_trace for ref in span.references)]
|
||||
return roots[0] if len(roots) == 1 else None
|
||||
|
||||
|
||||
def _follows(trace: JaegerTrace, parent_trace_id: str, parent_span_id: str) -> bool:
|
||||
root = root_span(trace)
|
||||
return root is not None and any(
|
||||
ref.trace_id == parent_trace_id and ref.span_id == parent_span_id for ref in root.references
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CallTraces:
|
||||
hits: tuple[JaegerTrace, ...]
|
||||
linked: tuple[JaegerTrace, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Observation:
|
||||
traces: CallTraces
|
||||
missing: tuple[str, ...]
|
||||
unreachable: NetworkError | None
|
||||
|
||||
def settled(self, names: set[str], prefixes: set[str]) -> bool:
|
||||
hits: Final = self.traces.hits
|
||||
return self.unreachable is None and len(hits) == 1 and not self.missing and _settled(hits[0], names, prefixes)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OtelReader:
|
||||
query_url: str
|
||||
|
|
@ -119,36 +165,95 @@ class OtelReader:
|
|||
case failure:
|
||||
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
|
||||
|
||||
def _query_operation(self, operation: str, *, start: int) -> Result[JaegerTracesPage]:
|
||||
return get(
|
||||
URL(f"{self.query_url}/api/traces"),
|
||||
headers=NoBody(),
|
||||
params=_TracesQuery(
|
||||
service=JAEGER_SERVICE,
|
||||
operation=operation,
|
||||
limit=200,
|
||||
start=start,
|
||||
end=int(time.time() * 1_000_000),
|
||||
),
|
||||
response_type=JaegerTracesPage,
|
||||
timeout=30.0,
|
||||
)
|
||||
|
||||
def linked_traces(self, *, operation: str, parent: JaegerTrace) -> tuple[JaegerTrace, ...] | NetworkError:
|
||||
"""Traces whose root span references the parent trace's root span.
|
||||
Detached post-response work lands as the root of its own trace with a
|
||||
link back to the request span instead of the call-id tag, so it is
|
||||
found by operation name, windowed to start at the parent root's start
|
||||
time (the detached span always starts after it), and matched on that
|
||||
link. A NetworkError is handed back so the polling caller can tell an
|
||||
unreachable read-back endpoint from a span that never arrived."""
|
||||
parent_root: Final = root_span(parent)
|
||||
if parent_root is None:
|
||||
return ()
|
||||
match self._query_operation(operation, start=parent_root.start_time):
|
||||
case Success(data=page):
|
||||
return tuple(t for t in page.data if _follows(t, parent.trace_id, parent_root.span_id))
|
||||
case NetworkError() as failure:
|
||||
return failure
|
||||
case failure:
|
||||
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
|
||||
|
||||
def poll_traces_for_call(
|
||||
self, *, call_id: str, settled_names: set[str], settled_prefixes: set[str]
|
||||
) -> list[JaegerTrace]:
|
||||
"""Poll until exactly one trace holds the call and it carries every span
|
||||
self,
|
||||
*,
|
||||
call_id: str,
|
||||
settled_names: set[str],
|
||||
settled_prefixes: set[str],
|
||||
linked_names: frozenset[str] = frozenset(),
|
||||
) -> CallTraces:
|
||||
"""Poll until exactly one trace holds the call, it carries every span
|
||||
name in ``settled_names`` plus at least one name per prefix in
|
||||
``settled_prefixes`` (spans flush in batches, the cost write lands after
|
||||
the response), then return the hits. At the deadline the last hits are
|
||||
returned as-is so the caller's assertions report the real final state -
|
||||
on a split trace this never settles and the orphan comes back."""
|
||||
deadline = time.monotonic() + POLL_TIMEOUT
|
||||
hits: list[JaegerTrace] = []
|
||||
unreachable: NetworkError | None = None
|
||||
while time.monotonic() < deadline:
|
||||
match self._query_traces(call_id):
|
||||
case Success(data=page):
|
||||
unreachable = None
|
||||
hits = page.data
|
||||
if len(hits) == 1 and _settled(hits[0], settled_names, settled_prefixes):
|
||||
return hits
|
||||
case NetworkError() as failure:
|
||||
unreachable = failure
|
||||
case failure:
|
||||
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
|
||||
time.sleep(POLL_INTERVAL)
|
||||
if unreachable is not None:
|
||||
``settled_prefixes``, and every name in ``linked_names`` is either in
|
||||
that trace or is the root of its own trace referencing the request
|
||||
root (post-response work detaches per #42826). At the deadline the
|
||||
last observed state is returned as-is so the caller's assertions
|
||||
report the real final state - on a split trace this never settles and
|
||||
the orphan comes back. A read-back endpoint still failing at the
|
||||
deadline (either query) is a hard failure, not a missing span."""
|
||||
deadline: Final = time.monotonic() + POLL_TIMEOUT
|
||||
last: Final = self._poll(call_id, settled_names, settled_prefixes, linked_names, deadline)
|
||||
if last.unreachable is not None:
|
||||
pytest.fail(
|
||||
f"Jaeger query API at {self.query_url} stayed unreachable until the "
|
||||
f"{POLL_TIMEOUT}s poll deadline: {unreachable}"
|
||||
f"{POLL_TIMEOUT}s poll deadline: {last.unreachable}"
|
||||
)
|
||||
return hits
|
||||
return last.traces
|
||||
|
||||
def _observe(self, call_id: str, linked_names: frozenset[str]) -> _Observation:
|
||||
match self._query_traces(call_id):
|
||||
case NetworkError() as failure:
|
||||
return _Observation(CallTraces((), ()), tuple(linked_names), failure)
|
||||
case Success(data=page):
|
||||
if len(page.data) != 1:
|
||||
return _Observation(CallTraces(tuple(page.data), ()), tuple(linked_names), None)
|
||||
hit: Final = page.data[0]
|
||||
present: Final = frozenset(hit.span_names())
|
||||
results: Final = {
|
||||
name: self.linked_traces(operation=name, parent=hit) for name in linked_names if name not in present
|
||||
}
|
||||
unreachable: Final = next((r for r in results.values() if isinstance(r, NetworkError)), None)
|
||||
linked: Final = tuple(t for r in results.values() if not isinstance(r, NetworkError) for t in r)
|
||||
missing: Final = tuple(name for name, r in results.items() if isinstance(r, NetworkError) or not r)
|
||||
return _Observation(CallTraces((hit,), linked), missing, unreachable)
|
||||
case failure:
|
||||
pytest.fail(f"Jaeger query API at {self.query_url} failed: {failure}")
|
||||
|
||||
def _poll(
|
||||
self,
|
||||
call_id: str,
|
||||
names: set[str],
|
||||
prefixes: set[str],
|
||||
linked_names: frozenset[str],
|
||||
deadline: float,
|
||||
) -> _Observation:
|
||||
observations: Final = (self._observe(call_id, linked_names) for _ in _ticks())
|
||||
return next(o for o in observations if o.settled(names, prefixes) or time.monotonic() >= deadline)
|
||||
|
||||
|
||||
def build_otel_reader() -> OtelReader:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue