more tests

Author: sentrivana

sentry_sdk/traces.py

@@ -230,7 +230,6 @@ class StreamedSpan:
         "_scope",
         "_flags",
         "_context_manager_state",
-        "_profile",
         "_continuous_profile",
         "_baggage",
         "sample_rate",
@@ -296,7 +295,6 @@ def __init__(
             self._sample_rand = _generate_sample_rand(self.trace_id)
 
         self._flags: dict[str, bool] = {}
-        self._profile = None
         self._continuous_profile: "Optional[ContinuousProfile]" = None
 
         self._update_active_thread()
@@ -320,7 +318,7 @@ def __enter__(self) -> "StreamedSpan":
 
         if self.is_segment():
             sampling_context = {
-                "transaction_context": {
+                "transaction_context": {  # XXX[span-first]: rename?
                     "trace_id": self.trace_id,
                     "span_id": self.span_id,
                     "parent_span_id": self.parent_span_id,
@@ -341,9 +339,6 @@ def __exit__(
         self, ty: "Optional[Any]", value: "Optional[Any]", tb: "Optional[Any]"
     ) -> None:
         if self.is_segment():
-            if self._profile is not None:
-                self._profile.__exit__(ty, value, tb)
-
             if self._continuous_profile is not None:
                 self._continuous_profile.stop()
 
@@ -440,6 +435,12 @@ def get_attributes(self) -> "Attributes":
     def set_attribute(self, key: str, value: "AttributeValue") -> None:
         self.attributes[key] = format_attribute(value)
 
+    def remove_attribute(self, key: str) -> None:
+        try:
+            del self.attributes[key]
+        except KeyError:
+            pass
+
     def set_attributes(self, attributes: "Attributes") -> None:
         for key, value in attributes.items():
             self.set_attribute(key, value)

tests/tracing/test_span_streaming.py

@@ -6,6 +6,7 @@
 import pytest
 
 import sentry_sdk
+from sentry_sdk.traces import SegmentSource
 
 minimum_python_38 = pytest.mark.skipif(
     sys.version_info < (3, 8), reason="Asyncio tests need Python >= 3.8"
@@ -43,9 +44,9 @@ def test_start_span(sentry_init, capture_envelopes):
 
     events = capture_envelopes()
 
-    with sentry_sdk.traces.start_span(name="segment"):
-        with sentry_sdk.traces.start_span(name="child"):
-            ...
+    with sentry_sdk.traces.start_span(name="segment") as segment:
+        with sentry_sdk.traces.start_span(name="child") as child:
+            assert child.segment == segment
 
     sentry_sdk.get_client().flush()
     spans = envelopes_to_spans(events)
@@ -85,6 +86,7 @@ def test_start_span_no_context_manager(sentry_init, capture_envelopes):
     segment.start()
     child = sentry_sdk.traces.start_span(name="child")
     child.start()
+    assert child.segment == segment
     child.finish()
     segment.finish()
 
@@ -113,6 +115,37 @@ def test_start_span_no_context_manager(sentry_init, capture_envelopes):
     assert segment["status"] == "ok"
 
 
+def test_span_sampled_at_start(sentry_init, capture_envelopes):
+    # Test that if a span is created without the context manager, it is sampled
+    # at .start() time rather then creation time
+
+    def traces_sampler(sampling_context):
+        assert "delayed_attribute" in sampling_context["attributes"]
+        assert sampling_context["attributes"]["delayed_attribute"] == 12
+        return 1.0
+
+    sentry_init(
+        traces_sampler=traces_sampler,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    events = capture_envelopes()
+
+    segment = sentry_sdk.traces.start_span(name="segment")
+    segment.set_attribute("delayed_attribute", 12)
+    segment.start()
+    segment.finish()
+
+    sentry_sdk.get_client().flush()
+    spans = envelopes_to_spans(events)
+
+    assert len(spans) == 1
+    (segment,) = spans
+
+    assert segment["name"] == "segment"
+    assert segment["attributes"]["delayed_attribute"] == 12
+
+
 def test_start_span_attributes(sentry_init, capture_envelopes):
     sentry_init(
         traces_sample_rate=1.0,
@@ -165,6 +198,79 @@ def traces_sampler(sampling_context):
     assert span["attributes"]["my_attribute"] == "my_value"
 
 
+def test_span_attributes(sentry_init, capture_envelopes):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    events = capture_envelopes()
+
+    class Class:
+        pass
+
+    with sentry_sdk.traces.start_span(
+        name="segment", attributes={"attribute1": "value"}
+    ) as span:
+        assert span.get_attributes()["attribute1"] == "value"
+        span.set_attribute("attribute2", 47)
+        span.remove_attribute("attribute1")
+        span.set_attributes({"attribute3": 4.5, "attribute4": False})
+        assert "attribute1" not in span.get_attributes()
+        attributes = span.get_attributes()
+        assert attributes["attribute2"] == 47
+        assert attributes["attribute3"] == 4.5
+        assert attributes["attribute4"] is False
+
+    sentry_sdk.get_client().flush()
+    spans = envelopes_to_spans(events)
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "segment"
+    assert "attribute1" not in span["attributes"]
+    assert span["attributes"]["attribute2"] == 47
+    assert span["attributes"]["attribute3"] == 4.5
+    assert span["attributes"]["attribute4"] is False
+
+
+def test_span_attributes_serialize_early(sentry_init, capture_envelopes):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    events = capture_envelopes()
+
+    class Class:
+        pass
+
+    with sentry_sdk.traces.start_span(name="span") as span:
+        span.set_attributes(
+            {
+                # arrays of different types will be serialized
+                "attribute1": [123, "text"],
+                # so will custom class instances
+                "attribute2": Class(),
+            }
+        )
+        attributes = span.get_attributes()
+        assert isinstance(attributes["attribute1"], str)
+        assert attributes["attribute1"] == "[123, 'text']"
+        assert isinstance(attributes["attribute2"], str)
+        assert "Class" in attributes["attribute2"]
+
+    sentry_sdk.get_client().flush()
+    spans = envelopes_to_spans(events)
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["attributes"]["attribute1"] == "[123, 'text']"
+    assert "Class" in span["attributes"]["attribute2"]
+
+
 def test_traces_sampler_drops_span(sentry_init, capture_envelopes):
     def traces_sampler(sampling_context):
         assert "attributes" in sampling_context
@@ -432,7 +538,6 @@ def traced_function(): ...
 
 
 @minimum_python_38
-@pytest.mark.asyncio
 def test_trace_decorator_async(sentry_init, capture_envelopes):
     sentry_init(
         traces_sample_rate=1.0,
@@ -459,6 +564,50 @@ async def traced_function(): ...
     assert span["status"] == "ok"
 
 
+def test_set_span_op(sentry_init, capture_envelopes):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    events = capture_envelopes()
+
+    with sentry_sdk.traces.start_span(name="span") as span:
+        span.set_op("function")
+        assert span.get_attributes()["sentry.op"] == "function"
+
+    sentry_sdk.get_client().flush()
+    spans = envelopes_to_spans(events)
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "span"
+    assert span["attributes"]["sentry.op"] == "function"
+
+
+def test_set_span_source(sentry_init, capture_envelopes):
+    sentry_init(
+        traces_sample_rate=1.0,
+        _experiments={"trace_lifecycle": "stream"},
+    )
+
+    events = capture_envelopes()
+
+    with sentry_sdk.traces.start_span(name="span") as span:
+        span.set_source(SegmentSource.TASK)
+        assert span.get_attributes()["sentry.span.source"] == SegmentSource.TASK.value
+
+    sentry_sdk.get_client().flush()
+    spans = envelopes_to_spans(events)
+
+    assert len(spans) == 1
+    (span,) = spans
+
+    assert span["name"] == "span"
+    assert span["attributes"]["sentry.span.source"] == SegmentSource.TASK.value
+
+
 def test_transport_format(sentry_init, capture_envelopes):
     sentry_init(
         server_name="test-server",