Skip to content

Replace remaining uses of fetch_ordered_spans and fetch_traces for stronger tests #288

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Mar 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 148 additions & 23 deletions tests/test_agent_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

from .fake_model import FakeModel
from .test_responses import get_text_message
from .testing_processor import fetch_normalized_spans, fetch_traces
from .testing_processor import assert_no_traces, fetch_normalized_spans


@pytest.mark.asyncio
Expand Down Expand Up @@ -164,7 +164,7 @@ async def test_parent_disabled_trace_disabled_agent_trace():

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

assert fetch_normalized_spans() == snapshot([])
assert_no_traces()


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

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

assert fetch_normalized_spans() == snapshot([])
assert_no_traces()


@pytest.mark.asyncio
Expand All @@ -193,16 +193,29 @@ async def test_trace_config_works():
await Runner.run(
agent,
input="first_test",
run_config=RunConfig(workflow_name="Foo bar", group_id="123", trace_id="456"),
run_config=RunConfig(workflow_name="Foo bar", group_id="123", trace_id="trace_456"),
)

traces = fetch_traces()
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
export = traces[0].export()
assert export is not None, "Trace export should not be None"
assert export["workflow_name"] == "Foo bar"
assert export["group_id"] == "123"
assert export["id"] == "456"
assert fetch_normalized_spans(keep_trace_id=True) == snapshot(
[
{
"id": "trace_456",
"workflow_name": "Foo bar",
"group_id": "123",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent",
"handoffs": [],
"tools": [],
"output_type": "str",
},
}
],
}
]
)


@pytest.mark.asyncio
Expand Down Expand Up @@ -259,8 +272,24 @@ async def test_streaming_single_run_is_single_trace():
async for _ in x.stream_events():
pass

traces = fetch_traces()
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
assert fetch_normalized_spans() == snapshot(
[
{
"workflow_name": "Agent workflow",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent",
"handoffs": [],
"tools": [],
"output_type": "str",
},
}
],
}
]
)


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

traces = fetch_traces()
assert len(traces) == 2, f"Expected 2 traces, got {len(traces)}"
assert fetch_normalized_spans() == snapshot(
[
{
"workflow_name": "Agent workflow",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
}
],
},
{
"workflow_name": "Agent workflow",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
}
],
},
]
)


@pytest.mark.asyncio
Expand Down Expand Up @@ -317,8 +376,42 @@ async def test_wrapped_streaming_trace_is_single_trace():
async for _ in x.stream_events():
pass

traces = fetch_traces()
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
assert fetch_normalized_spans() == snapshot(
[
{
"workflow_name": "test_workflow",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
],
}
]
)


@pytest.mark.asyncio
Expand Down Expand Up @@ -347,8 +440,42 @@ async def test_wrapped_mixed_trace_is_single_trace():
async for _ in x.stream_events():
pass

traces = fetch_traces()
assert len(traces) == 1, f"Expected 1 trace, got {len(traces)}"
assert fetch_normalized_spans() == snapshot(
[
{
"workflow_name": "test_workflow",
"children": [
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
{
"type": "agent",
"data": {
"name": "test_agent_1",
"handoffs": [],
"tools": [],
"output_type": "str",
},
},
],
}
]
)


@pytest.mark.asyncio
Expand All @@ -370,8 +497,7 @@ async def test_parent_disabled_trace_disables_streaming_agent_trace():
async for _ in x.stream_events():
pass

traces = fetch_traces()
assert len(traces) == 0, f"Expected 0 traces, got {len(traces)}"
assert_no_traces()


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

traces = fetch_traces()
assert len(traces) == 0, f"Expected 0 traces, got {len(traces)}"
assert_no_traces()
20 changes: 8 additions & 12 deletions tests/test_responses_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from agents.tracing.span_data import ResponseSpanData
from tests import fake_model

from .testing_processor import fetch_normalized_spans, fetch_ordered_spans
from .testing_processor import assert_no_spans, fetch_normalized_spans, fetch_ordered_spans


class DummyTracing:
Expand Down Expand Up @@ -89,9 +89,8 @@ async def dummy_fetch_response(
[{"workflow_name": "test", "children": [{"type": "response"}]}]
)

spans = fetch_ordered_spans()
assert len(spans) == 1
assert spans[0].span_data.response is None
[span] = fetch_ordered_spans()
assert span.span_data.response is None


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

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

spans = fetch_ordered_spans()
assert len(spans) == 0
assert_no_spans()


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

spans = fetch_ordered_spans()
assert len(spans) == 1
assert isinstance(spans[0].span_data, ResponseSpanData)
assert spans[0].span_data.response is None
[span] = fetch_ordered_spans()
assert isinstance(span.span_data, ResponseSpanData)
assert span.span_data.response is None


@pytest.mark.allow_call_model_methods
Expand Down Expand Up @@ -226,5 +223,4 @@ async def __aiter__(self):

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

spans = fetch_ordered_spans()
assert len(spans) == 0
assert_no_spans()
Loading