Skip to content

Commit 090e79b

Browse files
rm-openairm-openai
authored
Replace remaining uses of fetch_ordered_spans and fetch_traces for stronger tests (#288)
Following #261
2 parents 1d3d1ce + dacbb9b commit 090e79b

File tree

5 files changed

+323
-183
lines changed

5 files changed

+323
-183
lines changed

Diff for: ‎tests/test_agent_tracing.py

+148-23
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 fetch_normalized_spans, fetch_traces
12+
from .testing_processor import assert_no_traces, fetch_normalized_spans
1313

1414

1515
@pytest.mark.asyncio
@@ -164,7 +164,7 @@ async def test_parent_disabled_trace_disabled_agent_trace():
164164

165165
await Runner.run(agent, input="first_test")
166166

167-
assert fetch_normalized_spans() == snapshot([])
167+
assert_no_traces()
168168

169169

170170
@pytest.mark.asyncio
@@ -178,7 +178,7 @@ async def test_manual_disabling_works():
178178

179179
await Runner.run(agent, input="first_test", run_config=RunConfig(tracing_disabled=True))
180180

181-
assert fetch_normalized_spans() == snapshot([])
181+
assert_no_traces()
182182

183183

184184
@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
@@ -370,8 +497,7 @@ async def test_parent_disabled_trace_disables_streaming_agent_trace():
370497
async for _ in x.stream_events():
371498
pass
372499

373-
traces = fetch_traces()
374-
assert len(traces) == 0, f"Expected 0 traces, got {len(traces)}"
500+
assert_no_traces()
375501

376502

377503
@pytest.mark.asyncio
@@ -392,5 +518,4 @@ async def test_manual_streaming_disabling_works():
392518
async for _ in x.stream_events():
393519
pass
394520

395-
traces = fetch_traces()
396-
assert len(traces) == 0, f"Expected 0 traces, got {len(traces)}"
521+
assert_no_traces()

Diff for: ‎tests/test_responses_tracing.py

+8-12
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
10+
from .testing_processor import assert_no_spans, fetch_normalized_spans, fetch_ordered_spans
1111

1212

1313
class DummyTracing:
@@ -89,9 +89,8 @@ async def dummy_fetch_response(
8989
[{"workflow_name": "test", "children": [{"type": "response"}]}]
9090
)
9191

92-
spans = fetch_ordered_spans()
93-
assert len(spans) == 1
94-
assert spans[0].span_data.response is None
92+
[span] = fetch_ordered_spans()
93+
assert span.span_data.response is None
9594

9695

9796
@pytest.mark.allow_call_model_methods
@@ -116,8 +115,7 @@ async def dummy_fetch_response(
116115

117116
assert fetch_normalized_spans() == snapshot([{"workflow_name": "test"}])
118117

119-
spans = fetch_ordered_spans()
120-
assert len(spans) == 0
118+
assert_no_spans()
121119

122120

123121
@pytest.mark.allow_call_model_methods
@@ -190,10 +188,9 @@ async def __aiter__(self):
190188
[{"workflow_name": "test", "children": [{"type": "response"}]}]
191189
)
192190

193-
spans = fetch_ordered_spans()
194-
assert len(spans) == 1
195-
assert isinstance(spans[0].span_data, ResponseSpanData)
196-
assert spans[0].span_data.response is None
191+
[span] = fetch_ordered_spans()
192+
assert isinstance(span.span_data, ResponseSpanData)
193+
assert span.span_data.response is None
197194

198195

199196
@pytest.mark.allow_call_model_methods
@@ -226,5 +223,4 @@ async def __aiter__(self):
226223

227224
assert fetch_normalized_spans() == snapshot([{"workflow_name": "test"}])
228225

229-
spans = fetch_ordered_spans()
230-
assert len(spans) == 0
226+
assert_no_spans()

0 commit comments

Comments
 (0)