test(client): Add tests for dropped span client reports

Author: szokeasaurusrex

tests/test_client.py

@@ -1,8 +1,10 @@
+import contextlib
 import os
 import json
 import subprocess
 import sys
 import time
+from collections import Counter, defaultdict
 from collections.abc import Mapping
 from textwrap import dedent
 from unittest import mock
@@ -1214,3 +1216,192 @@ def test_uwsgi_warnings(sentry_init, recwarn, opt, missing_flags):
                 assert flag in str(record.message)
         else:
             assert not recwarn
+
+
+class TestSpanClientReports:
+    """
+    Tests for client reports related to spans.
+    """
+
+    class LostEventCapturingTransport(sentry_sdk.Transport):
+        """
+        A transport that captures lost events.
+        """
+
+        def __init__(self):
+            self.record_lost_event_calls = []
+            self.record_lost_transaction_calls = []
+
+        def capture_envelope(self, _):
+            pass
+
+        def record_lost_event(
+            self,
+            reason,
+            data_category=None,
+            item=None,
+            *,
+            quantity=1,
+        ):
+            self.record_lost_event_calls.append((reason, data_category, item, quantity))
+
+        def record_lost_transaction(
+            self,
+            reason,  # type: str
+            span_count,  # type: int
+        ):  # type: (...) -> None
+            self.record_lost_transaction_calls.append((reason, span_count))
+
+    @staticmethod
+    @contextlib.contextmanager
+    def patch_transport():
+        """Patches the transport with a new LostEventCapturingTransport, which we yield."""
+        old_transport = sentry_sdk.get_client().transport
+        new_transport = TestSpanClientReports.LostEventCapturingTransport()
+        sentry_sdk.get_client().transport = new_transport
+
+        try:
+            yield new_transport
+        finally:
+            sentry_sdk.get_client().transport = old_transport
+
+    @staticmethod
+    def span_dropper(spans_to_drop):
+        """
+        Returns a function that can be used to drop spans from an event.
+        """
+
+        def drop_spans(event, _):
+            event["spans"] = event["spans"][spans_to_drop:]
+            return event
+
+        return drop_spans
+
+    @staticmethod
+    def mock_transaction_event(span_count):
+        """
+        Returns a mock transaction event with the given number of spans.
+        """
+
+        return defaultdict(
+            mock.MagicMock,
+            type="transaction",
+            spans=[mock.MagicMock() for _ in range(span_count)],
+        )
+
+    def __init__(self, span_count):
+        """Configures a test case with the number of spans dropped and whether the transaction was dropped."""
+        self.span_count = span_count
+        self.expected_record_lost_event_calls = Counter()
+        self.expected_record_lost_transaction_calls = Counter()
+        self.before_send = lambda event, _: event
+        self.event_processor = lambda event, _: event
+        self.already_dropped_spans = 0
+
+    def _update_resulting_calls(
+        self, reason, drops_transaction=False, drops_spans=None
+    ):
+        """
+        Updates the expected calls with the given resulting calls.
+        """
+        if drops_transaction:
+            dropped_spans = self.span_count - self.already_dropped_spans
+            self.expected_record_lost_transaction_calls[(reason, dropped_spans)] += 1
+
+        elif drops_spans is not None:
+            self.already_dropped_spans += drops_spans
+            self.expected_record_lost_event_calls[
+                (reason, "span", None, drops_spans)
+            ] += 1
+
+    def with_before_send(
+        self,
+        before_send,
+        *,
+        drops_transaction=False,
+        drops_spans=None,
+    ):
+        """drops_transaction and drops_spans are mutually exclusive."""
+        self.before_send = before_send
+        self._update_resulting_calls(
+            "before_send",
+            drops_transaction,
+            drops_spans,
+        )
+
+        return self
+
+    def with_event_processor(
+        self,
+        event_processor,
+        *,
+        drops_transaction=False,
+        drops_spans=None,
+    ):
+        self.event_processor = event_processor
+        self._update_resulting_calls(
+            "event_processor",
+            drops_transaction,
+            drops_spans,
+        )
+
+        return self
+
+    def run(self):
+        """Runs the test case with the configured parameters."""
+        sentry_sdk.init(before_send_transaction=self.before_send)
+
+        with sentry_sdk.isolation_scope() as scope:
+            scope.add_event_processor(self.event_processor)
+            with self.patch_transport() as transport:
+                event = self.mock_transaction_event(self.span_count)
+                sentry_sdk.get_client().capture_event(event, scope=scope)
+
+        # We use counters to ensure that the calls are made the expected number of times, disregarding order.
+        assert (
+            Counter(transport.record_lost_event_calls)
+            == self.expected_record_lost_event_calls
+        )
+        assert (
+            Counter(transport.record_lost_transaction_calls)
+            == self.expected_record_lost_transaction_calls
+        )
+
+
+@pytest.mark.parametrize(
+    "test_config",
+    (
+        TestSpanClientReports(10),  # No spans dropped
+        TestSpanClientReports(0).with_before_send(
+            lambda e, _: None, drops_transaction=True
+        ),
+        TestSpanClientReports(10).with_before_send(
+            lambda e, _: None, drops_transaction=True
+        ),
+        TestSpanClientReports(10).with_before_send(
+            TestSpanClientReports.span_dropper(3), drops_spans=3
+        ),
+        TestSpanClientReports(10).with_before_send(
+            TestSpanClientReports.span_dropper(10), drops_spans=10
+        ),
+        TestSpanClientReports(10).with_event_processor(
+            lambda e, _: None, drops_transaction=True
+        ),
+        TestSpanClientReports(10).with_event_processor(
+            TestSpanClientReports.span_dropper(3), drops_spans=3
+        ),
+        TestSpanClientReports(10).with_event_processor(
+            TestSpanClientReports.span_dropper(10), drops_spans=10
+        ),
+        TestSpanClientReports(10)
+        .with_event_processor(TestSpanClientReports.span_dropper(3), drops_spans=3)
+        .with_before_send(TestSpanClientReports.span_dropper(5), drops_spans=5),
+        TestSpanClientReports(10)
+        .with_event_processor(TestSpanClientReports.span_dropper(3), drops_spans=3)
+        .with_before_send(
+            lambda e, _: None, drops_transaction=True
+        ),  # Test proper number of spans with each reason
+    ),
+)
+def test_dropped_transaction(test_config):
+    test_config.run()