Skip to content

Commit dacbb9b

Browse files
alexmojakialexmojaki
committed
More fetch_normalized_spans
1 parent 6b509e3 commit dacbb9b

File tree

4 files changed

+154
-24
lines changed

4 files changed

+154
-24
lines changed

Diff for: ‎tests/test_agent_tracing.py

+144-17
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
from .fake_model import FakeModel
1111
from .test_responses import get_text_message
12-
from .testing_processor import assert_no_traces, fetch_normalized_spans, fetch_traces
12+
from .testing_processor import assert_no_traces, fetch_normalized_spans
1313

1414

1515
@pytest.mark.asyncio
@@ -193,16 +193,29 @@ async def test_trace_config_works():
193193
await Runner.run(
194194
agent,
195195
input="first_test",
196-
run_config=RunConfig(workflow_name="Foo bar", group_id="123", trace_id="456"),
196+
run_config=RunConfig(workflow_name="Foo bar", group_id="123", trace_id="trace_456"),
197197
)
198198

199-
traces = fetch_traces()
200-
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
201-
export = traces[0].export()
202-
assert export is not None, "Trace export should not be None"
203-
assert export["workflow_name"] == "Foo bar"
204-
assert export["group_id"] == "123"
205-
assert export["id"] == "456"
199+
assert fetch_normalized_spans(keep_trace_id=True) == snapshot(
200+
[
201+
{
202+
"id": "trace_456",
203+
"workflow_name": "Foo bar",
204+
"group_id": "123",
205+
"children": [
206+
{
207+
"type": "agent",
208+
"data": {
209+
"name": "test_agent",
210+
"handoffs": [],
211+
"tools": [],
212+
"output_type": "str",
213+
},
214+
}
215+
],
216+
}
217+
]
218+
)
206219

207220

208221
@pytest.mark.asyncio
@@ -259,8 +272,24 @@ async def test_streaming_single_run_is_single_trace():
259272
async for _ in x.stream_events():
260273
pass
261274

262-
traces = fetch_traces()
263-
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
275+
assert fetch_normalized_spans() == snapshot(
276+
[
277+
{
278+
"workflow_name": "Agent workflow",
279+
"children": [
280+
{
281+
"type": "agent",
282+
"data": {
283+
"name": "test_agent",
284+
"handoffs": [],
285+
"tools": [],
286+
"output_type": "str",
287+
},
288+
}
289+
],
290+
}
291+
]
292+
)
264293

265294

266295
@pytest.mark.asyncio
@@ -285,8 +314,38 @@ async def test_multiple_streamed_runs_are_multiple_traces():
285314
async for _ in x.stream_events():
286315
pass
287316

288-
traces = fetch_traces()
289-
assert len(traces) == 2, f"Expected 2 traces, got {len(traces)}"
317+
assert fetch_normalized_spans() == snapshot(
318+
[
319+
{
320+
"workflow_name": "Agent workflow",
321+
"children": [
322+
{
323+
"type": "agent",
324+
"data": {
325+
"name": "test_agent_1",
326+
"handoffs": [],
327+
"tools": [],
328+
"output_type": "str",
329+
},
330+
}
331+
],
332+
},
333+
{
334+
"workflow_name": "Agent workflow",
335+
"children": [
336+
{
337+
"type": "agent",
338+
"data": {
339+
"name": "test_agent_1",
340+
"handoffs": [],
341+
"tools": [],
342+
"output_type": "str",
343+
},
344+
}
345+
],
346+
},
347+
]
348+
)
290349

291350

292351
@pytest.mark.asyncio
@@ -317,8 +376,42 @@ async def test_wrapped_streaming_trace_is_single_trace():
317376
async for _ in x.stream_events():
318377
pass
319378

320-
traces = fetch_traces()
321-
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
379+
assert fetch_normalized_spans() == snapshot(
380+
[
381+
{
382+
"workflow_name": "test_workflow",
383+
"children": [
384+
{
385+
"type": "agent",
386+
"data": {
387+
"name": "test_agent_1",
388+
"handoffs": [],
389+
"tools": [],
390+
"output_type": "str",
391+
},
392+
},
393+
{
394+
"type": "agent",
395+
"data": {
396+
"name": "test_agent_1",
397+
"handoffs": [],
398+
"tools": [],
399+
"output_type": "str",
400+
},
401+
},
402+
{
403+
"type": "agent",
404+
"data": {
405+
"name": "test_agent_1",
406+
"handoffs": [],
407+
"tools": [],
408+
"output_type": "str",
409+
},
410+
},
411+
],
412+
}
413+
]
414+
)
322415

323416

324417
@pytest.mark.asyncio
@@ -347,8 +440,42 @@ async def test_wrapped_mixed_trace_is_single_trace():
347440
async for _ in x.stream_events():
348441
pass
349442

350-
traces = fetch_traces()
351-
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
443+
assert fetch_normalized_spans() == snapshot(
444+
[
445+
{
446+
"workflow_name": "test_workflow",
447+
"children": [
448+
{
449+
"type": "agent",
450+
"data": {
451+
"name": "test_agent_1",
452+
"handoffs": [],
453+
"tools": [],
454+
"output_type": "str",
455+
},
456+
},
457+
{
458+
"type": "agent",
459+
"data": {
460+
"name": "test_agent_1",
461+
"handoffs": [],
462+
"tools": [],
463+
"output_type": "str",
464+
},
465+
},
466+
{
467+
"type": "agent",
468+
"data": {
469+
"name": "test_agent_1",
470+
"handoffs": [],
471+
"tools": [],
472+
"output_type": "str",
473+
},
474+
},
475+
],
476+
}
477+
]
478+
)
352479

353480

354481
@pytest.mark.asyncio

Diff for: ‎tests/test_responses_tracing.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from agents.tracing.span_data import ResponseSpanData
88
from tests import fake_model
99

10-
from .testing_processor import fetch_normalized_spans, fetch_ordered_spans, assert_no_spans
10+
from .testing_processor import assert_no_spans, fetch_normalized_spans, fetch_ordered_spans
1111

1212

1313
class DummyTracing:
@@ -117,6 +117,7 @@ async def dummy_fetch_response(
117117

118118
assert_no_spans()
119119

120+
120121
@pytest.mark.allow_call_model_methods
121122
@pytest.mark.asyncio
122123
async def test_stream_response_creates_trace(monkeypatch):
@@ -222,4 +223,4 @@ async def __aiter__(self):
222223

223224
assert fetch_normalized_spans() == snapshot([{"workflow_name": "test"}])
224225

225-
assert_no_spans()
226+
assert_no_spans()

Diff for: ‎tests/test_tracing.py

+1-3
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,9 @@
2020

2121
from .testing_processor import (
2222
SPAN_PROCESSOR_TESTING,
23+
assert_no_traces,
2324
fetch_events,
2425
fetch_normalized_spans,
25-
fetch_ordered_spans,
26-
fetch_traces,
27-
assert_no_traces,
2826
)
2927

3028
### HELPERS

Diff for: ‎tests/testing_processor.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -93,14 +93,18 @@ def assert_no_traces():
9393
assert_no_spans()
9494

9595

96-
def fetch_normalized_spans(keep_span_id: bool = False):
96+
def fetch_normalized_spans(
97+
keep_span_id: bool = False, keep_trace_id: bool = False
98+
) -> list[dict[str, Any]]:
9799
nodes: dict[tuple[str, str | None], dict[str, Any]] = {}
98100
traces = []
99101
for trace_obj in fetch_traces():
100102
trace = trace_obj.export()
101103
assert trace
102104
assert trace.pop("object") == "trace"
103-
assert trace.pop("id").startswith("trace_")
105+
assert trace["id"].startswith("trace_")
106+
if not keep_trace_id:
107+
del trace["id"]
104108
trace = {k: v for k, v in trace.items() if v is not None}
105109
nodes[(trace_obj.trace_id, None)] = trace
106110
traces.append(trace)

0 commit comments

Comments
 (0)