From 937964c80c1c83b60f8b298b0de2dcce16ba05f7 Mon Sep 17 00:00:00 2001
From: Neel Shah <neel.shah@sentry.io>
Date: Fri, 15 Nov 2024 17:01:04 +0100
Subject: [PATCH] Add sample_rate property and fix test_monitor

---
 sentry_sdk/tracing.py | 13 +++++++++++++
 tests/test_monitor.py | 12 +++++-------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
index 59971e274e..9c87a45903 100644
--- a/sentry_sdk/tracing.py
+++ b/sentry_sdk/tracing.py
@@ -1403,6 +1403,19 @@ def sampled(self):
         # type: () -> Optional[bool]
         return self._otel_span.get_span_context().trace_flags.sampled
 
+    @property
+    def sample_rate(self):
+        # type: () -> Optional[float]
+        from sentry_sdk.integrations.opentelemetry.consts import (
+            TRACESTATE_SAMPLE_RATE_KEY,
+        )
+
+        sample_rate = self._otel_span.get_span_context().trace_state.get(
+            TRACESTATE_SAMPLE_RATE_KEY
+        )
+        sample_rate = cast("Optional[str]", sample_rate)
+        return float(sample_rate) if sample_rate is not None else None
+
     @property
     def op(self):
         # type: () -> Optional[str]
diff --git a/tests/test_monitor.py b/tests/test_monitor.py
index 041169d515..1015e9f2ad 100644
--- a/tests/test_monitor.py
+++ b/tests/test_monitor.py
@@ -55,7 +55,7 @@ def test_monitor_unhealthy(sentry_init):
         assert monitor.downsample_factor == (i + 1 if i < 10 else 10)
 
 
-def test_transaction_uses_downsample_rate(
+def test_root_span_uses_downsample_rate(
     sentry_init, capture_envelopes, capture_record_lost_event_calls, monkeypatch
 ):
     sentry_init(
@@ -78,16 +78,14 @@ def test_transaction_uses_downsample_rate(
     assert monitor.is_healthy() is False
     assert monitor.downsample_factor == 1
 
-    with sentry_sdk.start_transaction(name="foobar") as transaction:
+    with sentry_sdk.start_span(name="foobar") as root_span:
         with sentry_sdk.start_span(name="foospan"):
             with sentry_sdk.start_span(name="foospan2"):
                 with sentry_sdk.start_span(name="foospan3"):
                     ...
 
-        assert transaction.sampled is False
-        assert (
-            transaction.sample_rate == 0.5
-        )  # TODO: this fails until we put the sample_rate in the POTelSpan
+        assert root_span.sampled is False
+        assert root_span.sample_rate == 0.5
 
     assert len(envelopes) == 0
 
@@ -104,7 +102,7 @@ def test_transaction_uses_downsample_rate(
                 "span",
                 None,
                 1,
-            ),  # Only one span (the transaction itself) is counted, since we did not record any spans in the first place.
+            ),  # Only one span (the root span itself) is counted, since we did not record any spans in the first place.
         ]
     )