tests: Move streaming decorator tests

tests/tracing/test_decorator.py

@@ -5,7 +5,6 @@
 
 import sentry_sdk
 from sentry_sdk.consts import SPANTEMPLATE
-from sentry_sdk.tracing import trace
 from sentry_sdk.tracing_utils import create_span_decorator
 from sentry_sdk.utils import logger
 from tests.conftest import patch_start_tracing_child
@@ -86,6 +85,246 @@ async def test_trace_decorator_async_no_trx():
             assert result2 == "return_of_async_function"
 
 
+def test_trace_decorator_span_streaming(sentry_init, capture_items):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace
+    def traced_function():
+        return "ok"
+
+    result = traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert (
+        span["name"]
+        == "test_decorator.test_trace_decorator_span_streaming.<locals>.traced_function"
+    )
+    assert span["status"] == "ok"
+
+
+def test_trace_decorator_arguments_span_streaming(sentry_init, capture_items):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace(name="traced", attributes={"traced.attribute": 123})
+    def traced_function():
+        return "ok"
+
+    result = traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "traced"
+    assert span["attributes"]["traced.attribute"] == 123
+    assert span["status"] == "ok"
+
+
+def test_trace_decorator_inactive_span_streaming(sentry_init, capture_items):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace(name="outer", active=False)
+    def traced_function():
+        with sentry_sdk.traces.start_span(name="inner"):
+            return "ok"
+
+    result = traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 2
+    (span1, span2) = spans
+
+    assert span1["name"] == "inner"
+    assert span1.get("parent_span_id") != span2["span_id"]
+
+    assert span2["name"] == "outer"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_span_streaming(sentry_init, capture_items):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace
+    async def traced_function():
+        return "ok"
+
+    result = await traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert (
+        span["name"]
+        == "test_decorator.test_trace_decorator_async_span_streaming.<locals>.traced_function"
+    )
+    assert span["status"] == "ok"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_arguments_span_streaming(
+    sentry_init, capture_items
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace(name="traced", attributes={"traced.attribute": 123})
+    async def traced_function():
+        return "ok"
+
+    result = await traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "traced"
+    assert span["attributes"]["traced.attribute"] == 123
+    assert span["status"] == "ok"
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_inactive_span_streaming(
+    sentry_init, capture_items
+):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace(name="outer", active=False)
+    async def traced_function():
+        with sentry_sdk.traces.start_span(name="inner"):
+            return "ok"
+
+    result = await traced_function()
+    assert result == "ok"
+
+    sentry_sdk.get_client().flush()
+    spans = [item.payload for item in items]
+
+    assert len(spans) == 2
+    (span1, span2) = spans
+
+    assert span1["name"] == "inner"
+    assert span1.get("parent_span_id") != span2["span_id"]
+
+    assert span2["name"] == "outer"
+
+
+def test_trace_decorator_child_span_streaming(sentry_init, capture_items):
+    """Spans created with @trace show up as children if a span is active."""
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+        },
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace
+    def _some_function_traced_stream(a, b, c):
+        return True
+
+    with sentry_sdk.traces.start_span(name="segment") as segment:
+        result = _some_function_traced_stream(1, 2, 3)
+
+    assert result is True
+
+    sentry_sdk.flush()
+
+    assert len(items) == 2
+    child_span, segment_span = items[0].payload, items[1].payload
+
+    assert (
+        child_span["name"]
+        == "test_decorator.test_trace_decorator_child_span_streaming.<locals>._some_function_traced_stream"
+    )
+    assert child_span["parent_span_id"] == segment.span_id
+    assert segment_span["name"] == "segment"
+    assert "parent_span_id" not in segment_span
+
+
+@pytest.mark.asyncio
+async def test_trace_decorator_async_child_span_streaming(sentry_init, capture_items):
+    """Spans created with @trace show up as children if a span is active."""
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={
+            "trace_lifecycle": "stream",
+        },
+    )
+
+    items = capture_items("span")
+
+    @sentry_sdk.traces.trace
+    async def _some_function_traced_stream(a, b, c):
+        return True
+
+    with sentry_sdk.traces.start_span(name="segment") as segment:
+        result = await _some_function_traced_stream(1, 2, 3)
+
+    assert result is True
+
+    sentry_sdk.flush()
+
+    assert len(items) == 2
+    child_span, segment_span = items[0].payload, items[1].payload
+
+    assert (
+        child_span["name"]
+        == "test_decorator.test_trace_decorator_async_child_span_streaming.<locals>._some_function_traced_stream"
+    )
+    assert child_span["parent_span_id"] == segment.span_id
+    assert segment_span["name"] == "segment"
+    assert "parent_span_id" not in segment_span
+
+
 def test_functions_to_trace_signature_unchanged_sync(sentry_init):
     sentry_init(
         traces_sample_rate=1.0,
@@ -94,14 +333,22 @@ def test_functions_to_trace_signature_unchanged_sync(sentry_init):
     def _some_function(a, b, c):
         pass
 
-    @trace
+    @sentry_sdk.trace
     def _some_function_traced(a, b, c):
         pass
 
+    @sentry_sdk.traces.trace
+    def _some_function_traced_stream(a, b, c):
+        pass
+
     assert inspect.getcallargs(_some_function, 1, 2, 3) == inspect.getcallargs(
         _some_function_traced, 1, 2, 3
     )
 
+    assert inspect.getcallargs(_some_function, 1, 2, 3) == inspect.getcallargs(
+        _some_function_traced_stream, 1, 2, 3
+    )
+
 
 @pytest.mark.asyncio
 async def test_functions_to_trace_signature_unchanged_async(sentry_init):
@@ -112,13 +359,20 @@ async def test_functions_to_trace_signature_unchanged_async(sentry_init):
     async def _some_function(a, b, c):
         pass
 
-    @trace
+    @sentry_sdk.trace
     async def _some_function_traced(a, b, c):
         pass
 
+    @sentry_sdk.traces.trace
+    async def _some_function_traced_stream(a, b, c):
+        pass
+
     assert inspect.getcallargs(_some_function, 1, 2, 3) == inspect.getcallargs(
         _some_function_traced, 1, 2, 3
     )
+    assert inspect.getcallargs(_some_function, 1, 2, 3) == inspect.getcallargs(
+        _some_function_traced_stream, 1, 2, 3
+    )
 
 
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])