chore: change place of assert_span_event_count/_attributes

DataDog · Feb 5, 2025 · 8202b1b · 8202b1b
1 parent 349990f
commit 8202b1b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/tests/tracer/test_tracer.py b/tests/tracer/test_tracer.py
@@ -179,8 +179,8 @@ def f():
                 error=0,
             ),
         )
-        self.assert_span_event_count(1)
-        self.assert_span_event_attributes(
+        self.spans[0].assert_span_event_count(1)
+        self.spans[0].assert_span_event_attributes(
             0, {"exception.type": "builtins.RuntimeError", "exception.message": "bim", "exception.escaped": "False"}
         )
 
@@ -198,17 +198,17 @@ def f():
                 self.tracer.record_exception(e)
 
         f()
-        self.assert_span_event_count(2)
         self.assert_structure(
             dict(
                 name="tests.tracer.test_tracer.f",
                 error=0,
             )
         )
-        self.assert_span_event_attributes(
+        self.spans[0].assert_span_event_count(2)
+        self.spans[0].assert_span_event_attributes(
             0, {"exception.type": "builtins.RuntimeError", "exception.message": "bim", "exception.escaped": "False"}
         )
-        self.assert_span_event_attributes(
+        self.spans[0].assert_span_event_attributes(
             1, {"exception.type": "builtins.RuntimeError", "exception.message": "bam", "exception.escaped": "False"}
         )
 
@@ -234,8 +234,8 @@ def f():
                 },
             ),
         )
-        self.assert_span_event_count(1)
-        self.assert_span_event_attributes(
+        self.spans[0].assert_span_event_count(1)
+        self.spans[0].assert_span_event_attributes(
             0, {"exception.type": "builtins.RuntimeError", "exception.message": "bim", "exception.escaped": "True"}
         )
 

diff --git a/tests/utils.py b/tests/utils.py
@@ -519,19 +519,6 @@ def assert_structure(self, root, children=NO_CHILDREN):
         root_span = self.get_root_span()
         root_span.assert_structure(root, children)
 
-    def assert_span_event_count(self, count):
-        "Assert the r"
-        root_span = self.get_root_span()
-        assert len(root_span._events) == count, "Span count {0} != {1}".format(len(root_span._events), count)
-
-    def assert_span_event_attributes(self, event_idx, attrs):
-        span_event_attrs = self.get_root_span()._events[event_idx].attributes
-        for name, value in attrs.items():
-            assert name in span_event_attrs, "{0!r} does not have property {1!r}".format(span_event_attrs, name)
-            assert span_event_attrs[name] == value, "{0!r} property {1}: {2!r} != {3!r}".format(
-                span_event_attrs, name, span_event_attrs[name], value
-            )
-
     @contextlib.contextmanager
     def override_global_tracer(self, tracer=None):
         original = ddtrace.tracer
@@ -860,6 +847,29 @@ def assert_metrics(self, metrics, exact=False):
                     self, key, self._metrics[key], value
                 )
 
+    def assert_span_event_count(self, count):
+        """Assert this span has the expected number of span_events"""
+        assert len(self._events) == count, "Span count {0} != {1}".format(len(self._events), count)
+
+    def assert_span_event_attributes(self, event_idx, attrs):
+        """
+        Assertion method to ensure this span's span event match as expected
+
+        Example::
+
+            span = TestSpan(span)
+            span.assert_span_event(0, {"exception.type": "builtins.RuntimeError"})
+
+        :param event_idx: id of the span event
+        :type event_idx: integer
+        """
+        span_event_attrs = self._events[event_idx].attributes
+        for name, value in attrs.items():
+            assert name in span_event_attrs, "{0!r} does not have property {1!r}".format(span_event_attrs, name)
+            assert span_event_attrs[name] == value, "{0!r} property {1}: {2!r} != {3!r}".format(
+                span_event_attrs, name, span_event_attrs[name], value
+            )
+
 
 class TracerSpanContainer(TestSpanContainer):
     """