test: Add test for custom sample rates

Author: szokeasaurusrex

tests/sentry/ingest/ingest_consumer/test_ingest_consumer_processing.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import datetime
+import os
 import time
 import uuid
 import zipfile
 from io import BytesIO
 from typing import Any
 from unittest.mock import Mock, patch
 
+import django.test
 import orjson
 import pytest
 from arroyo.backends.kafka.consumer import KafkaPayload
@@ -24,6 +26,7 @@
     process_event,
     process_individual_attachment,
     process_userreport,
+    trace_func,
 )
 from sentry.ingest.types import ConsumerType
 from sentry.models.debugfile import create_files_from_dif_zip
@@ -621,3 +624,58 @@ def test_collect_span_metrics(default_project):
             assert mock_metrics.incr.call_count == 0
             collect_span_metrics(default_project, {"spans": [1, 2, 3]})
             assert mock_metrics.incr.call_count == 1
+
+
+@pytest.mark.parametrize(
+    "env_value,settings_value,expected_sample_rate",
+    [
+        # Both unset - should use default of 0
+        (None, None, 0.0),
+        # Only environment variable set
+        ("0", None, 0.0),
+        ("1", None, 1.0),
+        ("0.5", None, 0.5),
+        # Only settings value set
+        (None, 0, 0.0),
+        (None, 1, 1.0),
+        (None, 0.7, 0.7),
+        # Both set - environment variable should take precedence
+        ("0", 1, 0.0),  # env=0, settings=1 -> should use env (0)
+        ("1", 0, 1.0),  # env=1, settings=0 -> should use env (1)
+        ("0.3", 0.8, 0.3),  # env=0.3, settings=0.8 -> should use env (0.3)
+    ],
+)
+def test_sample_rate_passed(env_value, settings_value, expected_sample_rate):
+    # Test various combinations of environment variable and settings values
+
+    # Prepare environment
+    env_dict = {}
+    if env_value is not None:
+        env_dict["SENTRY_INGEST_CONSUMER_APM_SAMPLING"] = env_value
+
+    with patch.dict(os.environ, env_dict, clear=True):
+        with django.test.override_settings(SENTRY_INGEST_CONSUMER_APM_SAMPLING=settings_value):
+            # If settings_value is None, delete the setting to simulate it not being set
+            if settings_value is None:
+                del settings.SENTRY_INGEST_CONSUMER_APM_SAMPLING
+
+            with patch(
+                "sentry.ingest.consumer.processors.sentry_sdk.start_span"
+            ) as mock_start_span:
+                # Create a dummy function to decorate
+                @trace_func(name="test_span")
+                def dummy_function():
+                    return "test_result"
+
+                # Call the decorated function
+                result = dummy_function()
+
+                # Verify the function returned correctly
+                assert result == "test_result"
+
+                # Verify start_span was called with correct arguments
+                mock_start_span.assert_called_once()
+                call_args = mock_start_span.call_args
+
+                # Check that the span_kwargs include the expected sample_rate
+                assert call_args.kwargs["attributes"]["sample_rate"] == expected_sample_rate

tests/sentry/utils/test_sdk.py

@@ -19,6 +19,7 @@
     check_current_scope_transaction,
     check_tag_for_scope_bleed,
     merge_context_into_scope,
+    traces_sampler,
 )
 
 
@@ -64,6 +65,18 @@ def test_context_scope_merge_with_existing_context(self):
             "charlie": "goofy",
         }
 
+    def test_traces_sampler_custom_sample_rate_0_0(self):
+        sampling_context = {"sample_rate": 0.0}
+        assert traces_sampler(sampling_context) == 0.0
+
+    def test_traces_sampler_custom_sample_rate_0_5(self):
+        sampling_context = {"sample_rate": 0.5}
+        assert traces_sampler(sampling_context) == 0.5
+
+    def test_traces_sampler_custom_sample_rate_1_0(self):
+        sampling_context = {"sample_rate": 1.0}
+        assert traces_sampler(sampling_context) == 1.0
+
 
 @patch("sentry.utils.sdk.logger.warning")
 class CheckTagForScopeBleedTest(TestCase):