tests(profiling): Add additional test coverage for profiler (getsentry#1877)

Zylphrex · web-flow · commit bac5bb1492d9 · 2023-01-31T08:08:55.000-05:00
tests(profiling): Add additional test coverage for profiler
diff --git a/sentry_sdk/profiler.py b/sentry_sdk/profiler.py
@@ -137,7 +137,7 @@ def is_gevent():
 
 
 def setup_profiler(options):
-    # type: (Dict[str, Any]) -> None
+    # type: (Dict[str, Any]) -> bool
     """
     `buffer_secs` determines the max time a sample will be buffered for
     `frequency` determines the number of samples to take per second (Hz)
@@ -147,11 +147,11 @@ def setup_profiler(options):
 
     if _scheduler is not None:
         logger.debug("profiling is already setup")
-        return
+        return False
 
     if not PY33:
         logger.warn("profiling is only supported on Python >= 3.3")
-        return
+        return False
 
     frequency = 101
 
@@ -184,6 +184,8 @@ def setup_profiler(options):
 
     atexit.register(teardown_profiler)
 
+    return True
+
 
 def teardown_profiler():
     # type: () -> None
@@ -410,8 +412,7 @@ def __init__(
         #
         # We cannot keep a reference to the transaction around here because it'll create
         # a reference cycle. So we opt to pull out just the necessary attributes.
-        self._transaction_sampled = transaction.sampled  # type: Optional[bool]
-        self.sampled = None  # type: Optional[bool]
+        self.sampled = transaction.sampled  # type: Optional[bool]
 
         # Various framework integrations are capable of overwriting the active thread id.
         # If it is set to `None` at the end of the profile, we fall back to the default.
@@ -448,7 +449,7 @@ def _set_initial_sampling_decision(self, sampling_context):
 
         # The corresponding transaction was not sampled,
         # so don't generate a profile for it.
-        if not self._transaction_sampled:
+        if not self.sampled:
             self.sampled = False
             return
 
@@ -485,19 +486,21 @@ def get_profile_context(self):
 
     def start(self):
         # type: () -> None
-        if not self.sampled:
+        if not self.sampled or self.active:
             return
 
         assert self.scheduler, "No scheduler specified"
+        self.active = True
         self.start_ns = nanosecond_time()
         self.scheduler.start_profiling(self)
 
     def stop(self):
         # type: () -> None
-        if not self.sampled:
+        if not self.sampled or not self.active:
             return
 
         assert self.scheduler, "No scheduler specified"
+        self.active = False
         self.scheduler.stop_profiling(self)
         self.stop_ns = nanosecond_time()
 
@@ -526,11 +529,15 @@ def __exit__(self, ty, value, tb):
 
     def write(self, ts, sample):
         # type: (int, RawSample) -> None
+        if not self.active:
+            return
+
         if ts < self.start_ns:
             return
 
         offset = ts - self.start_ns
         if offset > MAX_PROFILE_DURATION_NS:
+            self.stop()
             return
 
         elapsed_since_start_ns = str(offset)
@@ -666,12 +673,11 @@ def teardown(self):
 
     def start_profiling(self, profile):
         # type: (Profile) -> None
-        profile.active = True
         self.new_profiles.append(profile)
 
     def stop_profiling(self, profile):
         # type: (Profile) -> None
-        profile.active = False
+        pass
 
     def make_sampler(self):
         # type: () -> Callable[..., None]
diff --git a/tests/integrations/wsgi/test_wsgi.py b/tests/integrations/wsgi/test_wsgi.py
@@ -287,49 +287,15 @@ def sample_app(environ, start_response):
 @pytest.mark.skipif(
     sys.version_info < (3, 3), reason="Profiling is only supported in Python >= 3.3"
 )
-@pytest.mark.parametrize(
-    "profiles_sample_rate,profile_count",
-    [
-        pytest.param(1.0, 1, id="profiler sampled at 1.0"),
-        pytest.param(0.75, 1, id="profiler sampled at 0.75"),
-        pytest.param(0.25, 0, id="profiler not sampled at 0.25"),
-        pytest.param(None, 0, id="profiler not enabled"),
-    ],
-)
 def test_profile_sent(
     sentry_init,
     capture_envelopes,
     teardown_profiling,
-    profiles_sample_rate,
-    profile_count,
 ):
     def test_app(environ, start_response):
         start_response("200 OK", [])
         return ["Go get the ball! Good dog!"]
 
-    sentry_init(
-        traces_sample_rate=1.0,
-        _experiments={"profiles_sample_rate": profiles_sample_rate},
-    )
-    app = SentryWsgiMiddleware(test_app)
-    envelopes = capture_envelopes()
-
-    with mock.patch("sentry_sdk.profiler.random.random", return_value=0.5):
-        client = Client(app)
-        client.get("/")
-
-    count_item_types = Counter()
-    for envelope in envelopes:
-        for item in envelope.items:
-            count_item_types[item.type] += 1
-    assert count_item_types["profile"] == profile_count
-
-
-def test_profile_context_sent(sentry_init, capture_envelopes, teardown_profiling):
-    def test_app(environ, start_response):
-        start_response("200 OK", [])
-        return ["Go get the ball! Good dog!"]
-
     sentry_init(
         traces_sample_rate=1.0,
         _experiments={"profiles_sample_rate": 1.0},
@@ -340,19 +306,8 @@ def test_app(environ, start_response):
     client = Client(app)
     client.get("/")
 
-    transaction = None
-    profile = None
-    for envelope in envelopes:
-        for item in envelope.items:
-            if item.type == "profile":
-                assert profile is None  # should only have 1 profile
-                profile = item
-            elif item.type == "transaction":
-                assert transaction is None  # should only have 1 transaction
-                transaction = item
-
-    assert transaction is not None
-    assert profile is not None
-    assert transaction.payload.json["contexts"]["profile"] == {
-        "profile_id": profile.payload.json["event_id"],
-    }
+    envelopes = [envelope for envelope in envelopes]
+    assert len(envelopes) == 1
+
+    profiles = [item for item in envelopes[0].items if item.type == "profile"]
+    assert len(profiles) == 1
diff --git a/tests/test_profiler.py b/tests/test_profiler.py
