From cb0eb8e254dbb83259b9ca89de646da138032007 Mon Sep 17 00:00:00 2001 From: Alex Hall Date: Fri, 21 Mar 2025 18:09:19 +0200 Subject: [PATCH] More fetch_normalized_spans --- tests/test_tracing.py | 99 +++++++++++++++++++------------------- tests/testing_processor.py | 6 ++- uv.lock | 2 +- 3 files changed, 55 insertions(+), 52 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index c54c3d8..8368ff2 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -4,6 +4,7 @@ import asyncio from typing import Any import pytest +from inline_snapshot import snapshot from agents.tracing import ( Span, @@ -17,7 +18,13 @@ from agents.tracing import ( ) from agents.tracing.spans import SpanError -from .testing_processor import fetch_events, fetch_ordered_spans, fetch_traces +from .testing_processor import ( + SPAN_PROCESSOR_TESTING, + fetch_events, + fetch_normalized_spans, + fetch_ordered_spans, + fetch_traces, +) ### HELPERS @@ -129,11 +136,11 @@ def test_ctxmanager_spans() -> None: async def run_subtask(span_id: str | None = None) -> None: with generation_span(span_id=span_id): - await asyncio.sleep(0.01) + await asyncio.sleep(0.0001) async def simple_async_tracing(): - with trace(workflow_name="test", trace_id="123", group_id="456"): + with trace(workflow_name="test", trace_id="trace_123", group_id="group_456"): await run_subtask(span_id="span_1") await run_subtask(span_id="span_2") @@ -142,21 +149,18 @@ async def simple_async_tracing(): async def test_async_tracing() -> None: await simple_async_tracing() - spans, traces = fetch_ordered_spans(), fetch_traces() - assert len(spans) == 2 - assert len(traces) == 1 - - trace = traces[0] - standard_trace_checks(trace, name_check="test") - trace_id = trace.trace_id - - # We don't care about ordering here, just that they're there - for s in spans: - standard_span_checks(s, trace_id=trace_id, parent_id=None, span_type="generation") - - ids = [span.span_id for span in spans] - assert "span_1" in ids - assert "span_2" in ids + assert fetch_normalized_spans(keep_span_id=True) == snapshot( + [ + { + "workflow_name": "test", + "group_id": "group_456", + "children": [ + {"type": "generation", "id": "span_1"}, + {"type": "generation", "id": "span_2"}, + ], + } + ] + ) async def run_tasks_parallel(span_ids: list[str]) -> None: @@ -171,13 +175,11 @@ async def run_tasks_as_children(first_span_id: str, second_span_id: str) -> None async def complex_async_tracing(): - with trace(workflow_name="test", trace_id="123", group_id="456"): - await asyncio.sleep(0.01) + with trace(workflow_name="test", trace_id="trace_123", group_id="456"): await asyncio.gather( run_tasks_parallel(["span_1", "span_2"]), run_tasks_parallel(["span_3", "span_4"]), ) - await asyncio.sleep(0.01) await asyncio.gather( run_tasks_as_children("span_5", "span_6"), run_tasks_as_children("span_7", "span_8"), @@ -186,35 +188,34 @@ async def complex_async_tracing(): @pytest.mark.asyncio async def test_complex_async_tracing() -> None: - await complex_async_tracing() + for _ in range(300): + SPAN_PROCESSOR_TESTING.clear() + await complex_async_tracing() - spans, traces = fetch_ordered_spans(), fetch_traces() - assert len(spans) == 8 - assert len(traces) == 1 - - trace = traces[0] - standard_trace_checks(trace, name_check="test") - trace_id = trace.trace_id - - # First ensure 1,2,3,4 exist and are in parallel with the trace as parent - for span_id in ["span_1", "span_2", "span_3", "span_4"]: - span = next((s for s in spans if s.span_id == span_id), None) - assert span is not None - standard_span_checks(span, trace_id=trace_id, parent_id=None, span_type="generation") - - # Ensure 5 and 7 exist and have the trace as parent - for span_id in ["span_5", "span_7"]: - span = next((s for s in spans if s.span_id == span_id), None) - assert span is not None - standard_span_checks(span, trace_id=trace_id, parent_id=None, span_type="generation") - - # Ensure 6 and 8 exist and have 5 and 7 as parents - six = next((s for s in spans if s.span_id == "span_6"), None) - assert six is not None - standard_span_checks(six, trace_id=trace_id, parent_id="span_5", span_type="generation") - eight = next((s for s in spans if s.span_id == "span_8"), None) - assert eight is not None - standard_span_checks(eight, trace_id=trace_id, parent_id="span_7", span_type="generation") + assert fetch_normalized_spans(keep_span_id=True) == ( + [ + { + "workflow_name": "test", + "group_id": "456", + "children": [ + {"type": "generation", "id": "span_1"}, + {"type": "generation", "id": "span_2"}, + {"type": "generation", "id": "span_3"}, + {"type": "generation", "id": "span_4"}, + { + "type": "generation", + "id": "span_5", + "children": [{"type": "generation", "id": "span_6"}], + }, + { + "type": "generation", + "id": "span_7", + "children": [{"type": "generation", "id": "span_8"}], + }, + ], + } + ] + ) def spans_with_setters(): diff --git a/tests/testing_processor.py b/tests/testing_processor.py index 371ea86..b07dea5 100644 --- a/tests/testing_processor.py +++ b/tests/testing_processor.py @@ -80,7 +80,7 @@ def fetch_events() -> list[TestSpanProcessorEvent]: return SPAN_PROCESSOR_TESTING._events -def fetch_normalized_spans(): +def fetch_normalized_spans(keep_span_id: bool = False): nodes: dict[tuple[str, str | None], dict[str, Any]] = {} traces = [] for trace_obj in fetch_traces(): @@ -99,7 +99,9 @@ def fetch_normalized_spans(): span = span_obj.export() assert span assert span.pop("object") == "trace.span" - assert span.pop("id").startswith("span_") + assert span["id"].startswith("span_") + if not keep_span_id: + del span["id"] assert datetime.fromisoformat(span.pop("started_at")) assert datetime.fromisoformat(span.pop("ended_at")) parent_id = span.pop("parent_id") diff --git a/uv.lock b/uv.lock index 15e091a..a9c79e2 100644 --- a/uv.lock +++ b/uv.lock @@ -1050,7 +1050,7 @@ wheels = [ [[package]] name = "openai-agents" -version = "0.0.5" +version = "0.0.6" source = { editable = "." } dependencies = [ { name = "griffe" },