ref: Convert integration tests about dropping transactions to unit tests

relay-server/src/actors/processor.rs

@@ -2405,7 +2405,12 @@ impl Service for EnvelopeProcessorService {
 mod tests {
     use chrono::{DateTime, TimeZone, Utc};
     use relay_general::pii::{DataScrubbingConfig, PiiConfig};
+    use relay_general::protocol::EventId;
+    use relay_sampling::{RuleCondition, RuleType, SamplingMode};
 
+    use crate::service::ServiceState;
+    use crate::testutils::{new_envelope, state_with_rule_and_condition};
+    use crate::utils::Semaphore as TestSemaphore;
     use crate::{actors::project::ProjectConfig, extractors::RequestMeta};
 
     use super::*;
@@ -2440,6 +2445,57 @@ mod tests {
             .unwrap()
     }
 
+    #[test]
+    fn test_it_keeps_or_drops_transactions() {
+        relay_test::setup();
+
+        // an empty json still produces a valid config
+        let json_config = serde_json::json!({});
+
+        let config = Config::from_json_value(json_config.clone()).unwrap();
+        let arconfig = Arc::new(Config::from_json_value(json_config).unwrap());
+
+        // it really shouldn't be necessary to start an entire service for a unit test
+        // it was ported from a python integration test
+        ServiceState::start(arconfig).unwrap();
+        let service = create_test_processor(config);
+        let envelope = new_envelope(false, "foo");
+
+        let event = Event {
+            id: Annotated::new(EventId::new()),
+            ty: Annotated::new(EventType::Transaction),
+            transaction: Annotated::new("testing".to_owned()),
+            ..Event::default()
+        };
+
+        for (sample_rate, shouldkeep) in [(0.0, false), (1.0, true)] {
+            let project_state = state_with_rule_and_condition(
+                Some(sample_rate),
+                RuleType::Transaction,
+                SamplingMode::Received,
+                RuleCondition::all(),
+            );
+
+            let mut state = ProcessEnvelopeState {
+                envelope: new_envelope(false, "foo"),
+                event: Annotated::from(event.clone()),
+                transaction_metrics_extracted: false,
+                metrics: Default::default(),
+                sample_rates: None,
+                extracted_metrics: vec![],
+                project_state: Arc::new(project_state),
+                sampling_project_state: None,
+                project_id: ProjectId::new(42),
+                envelope_context: EnvelopeContext::new(
+                    &envelope,
+                    TestSemaphore::new(42).try_acquire().unwrap(),
+                ),
+            };
+            let result = service.sample_envelope(&mut state);
+            assert_eq!(result.is_ok(), shouldkeep);
+        }
+    }
+
     #[test]
     fn test_breadcrumbs_file1() {
         let item = create_breadcrumbs_item(&[(None, "item1")]);

relay-server/src/lib.rs

@@ -266,6 +266,9 @@ mod service;
 mod statsd;
 mod utils;
 
+#[cfg(test)]
+mod testutils;
+
 use relay_config::Config;
 use relay_system::Controller;
 

relay-server/src/testutils.rs

@@ -0,0 +1,115 @@
+use bytes::Bytes;
+use relay_general::protocol::EventId;
+use relay_sampling::{
+    DynamicSamplingContext, RuleCondition, RuleId, RuleType, SamplingConfig, SamplingMode,
+    SamplingRule,
+};
+
+use crate::actors::project::ProjectState;
+use crate::envelope::{Envelope, Item, ItemType};
+
+pub fn state_with_rule_and_condition(
+    sample_rate: Option<f64>,
+    rule_type: RuleType,
+    mode: SamplingMode,
+    condition: RuleCondition,
+) -> ProjectState {
+    let rules = match sample_rate {
+        Some(sample_rate) => vec![SamplingRule {
+            condition,
+            sample_rate,
+            ty: rule_type,
+            id: RuleId(1),
+            time_range: Default::default(),
+            decaying_fn: relay_sampling::DecayingFunction::Constant,
+        }],
+        None => Vec::new(),
+    };
+
+    state_with_config(SamplingConfig {
+        rules,
+        mode,
+        next_id: None,
+    })
+}
+
+pub fn state_with_config(sampling_config: SamplingConfig) -> ProjectState {
+    let mut state = ProjectState::allowed();
+    state.config.dynamic_sampling = Some(sampling_config);
+    state
+}
+
+pub fn state_with_rule(
+    sample_rate: Option<f64>,
+    rule_type: RuleType,
+    mode: SamplingMode,
+) -> ProjectState {
+    let rules = match sample_rate {
+        Some(sample_rate) => vec![SamplingRule {
+            condition: RuleCondition::all(),
+            sample_rate,
+            ty: rule_type,
+            id: RuleId(1),
+            time_range: Default::default(),
+            decaying_fn: relay_sampling::DecayingFunction::Constant,
+        }],
+        None => Vec::new(),
+    };
+
+    state_with_config(SamplingConfig {
+        rules,
+        mode,
+        next_id: None,
+    })
+}
+pub fn create_sampling_context(sample_rate: Option<f64>) -> DynamicSamplingContext {
+    DynamicSamplingContext {
+        trace_id: uuid::Uuid::new_v4(),
+        public_key: "12345678901234567890123456789012".parse().unwrap(),
+        release: None,
+        environment: None,
+        transaction: None,
+        sample_rate,
+        user: Default::default(),
+        other: Default::default(),
+    }
+}
+
+/// ugly hack to build an envelope with an optional trace context
+pub fn new_envelope<T: Into<String>>(with_dsc: bool, transaction_name: T) -> Box<Envelope> {
+    let transaction_name = transaction_name.into();
+    let dsn = "https://e12d836b15bb49d7bbf99e64295d995b:@sentry.io/42";
+    let event_id = EventId::new();
+
+    let raw_event = if with_dsc {
+        format!(
+            "{{\"transaction\": \"{}\", \"event_id\":\"{}\",\"dsn\":\"{}\", \"trace\": {}}}\n",
+            transaction_name,
+            event_id.0.to_simple(),
+            dsn,
+            serde_json::to_string(&create_sampling_context(None)).unwrap(),
+        )
+    } else {
+        format!(
+            "{{\"transaction\": \"{}\", \"event_id\":\"{}\",\"dsn\":\"{}\"}}\n",
+            transaction_name,
+            event_id.0.to_simple(),
+            dsn,
+        )
+    };
+
+    let bytes = Bytes::from(raw_event);
+
+    let mut envelope = Envelope::parse_bytes(bytes).unwrap();
+
+    let item1 = Item::new(ItemType::Transaction);
+    envelope.add_item(item1);
+
+    let item2 = Item::new(ItemType::Attachment);
+    envelope.add_item(item2);
+
+    let item3 = Item::new(ItemType::Attachment);
+    envelope.add_item(item3);
+
+    envelope
+}

relay-server/src/utils/dynamic_sampling.rs

@@ -184,7 +184,6 @@ pub fn get_sampling_key(envelope: &Envelope) -> Option<ProjectKey> {
 #[cfg(test)]
 mod tests {
 
-    use bytes::Bytes;
     use chrono::DateTime;
 
     use chrono::Duration as DateDuration;
@@ -196,31 +195,14 @@ mod tests {
         SamplingRule, TimeRange,
     };
 
-    use crate::envelope::Item;
+    use crate::testutils::create_sampling_context;
+    use crate::testutils::new_envelope;
+    use crate::testutils::state_with_config;
+    use crate::testutils::state_with_rule;
+    use crate::testutils::state_with_rule_and_condition;
 
     use super::*;
 
-    fn state_with_config(sampling_config: SamplingConfig) -> ProjectState {
-        let mut state = ProjectState::allowed();
-        state.config.dynamic_sampling = Some(sampling_config);
-        state
-    }
-
-    fn state_with_rule(
-        sample_rate: Option<f64>,
-        rule_type: RuleType,
-        mode: SamplingMode,
-    ) -> ProjectState {
-        state_with_decaying_rule(
-            sample_rate,
-            rule_type,
-            mode,
-            DecayingFunction::Constant,
-            None,
-            None,
-        )
-    }
-
     fn state_with_decaying_rule(
         sample_rate: Option<f64>,
         rule_type: RuleType,
@@ -248,58 +230,6 @@ mod tests {
         })
     }
 
-    fn create_sampling_context(sample_rate: Option<f64>) -> DynamicSamplingContext {
-        DynamicSamplingContext {
-            trace_id: uuid::Uuid::new_v4(),
-            public_key: "12345678901234567890123456789012".parse().unwrap(),
-            release: None,
-            environment: None,
-            transaction: None,
-            sample_rate,
-            user: Default::default(),
-            other: Default::default(),
-        }
-    }
-
-    /// ugly hack to build an envelope with an optional trace context
-    fn new_envelope<T: Into<String>>(with_dsc: bool, transaction_name: T) -> Box<Envelope> {
-        let transaction_name = transaction_name.into();
-        let dsn = "https://e12d836b15bb49d7bbf99e64295d995b:@sentry.io/42";
-        let event_id = EventId::new();
-
-        let raw_event = if with_dsc {
-            format!(
-                "{{\"transaction\": \"{}\", \"event_id\":\"{}\",\"dsn\":\"{}\", \"trace\": {}}}\n",
-                transaction_name,
-                event_id.0.to_simple(),
-                dsn,
-                serde_json::to_string(&create_sampling_context(None)).unwrap(),
-            )
-        } else {
-            format!(
-                "{{\"transaction\": \"{}\", \"event_id\":\"{}\",\"dsn\":\"{}\"}}\n",
-                transaction_name,
-                event_id.0.to_simple(),
-                dsn,
-            )
-        };
-
-        let bytes = Bytes::from(raw_event);
-
-        let mut envelope = Envelope::parse_bytes(bytes).unwrap();
-
-        let item1 = Item::new(ItemType::Transaction);
-        envelope.add_item(item1);
-
-        let item2 = Item::new(ItemType::Attachment);
-        envelope.add_item(item2);
-
-        let item3 = Item::new(ItemType::Attachment);
-        envelope.add_item(item3);
-
-        envelope
-    }
-
     fn prepare_and_get_sampling_rule(
         client_sample_rate: f64,
         event_type: EventType,
@@ -323,31 +253,6 @@ mod tests {
         )
     }
 
-    fn state_with_rule_and_condition(
-        sample_rate: Option<f64>,
-        rule_type: RuleType,
-        mode: SamplingMode,
-        condition: RuleCondition,
-    ) -> ProjectState {
-        let rules = match sample_rate {
-            Some(sample_rate) => vec![SamplingRule {
-                condition,
-                sample_rate,
-                ty: rule_type,
-                id: RuleId(1),
-                time_range: Default::default(),
-                decaying_fn: Default::default(),
-            }],
-            None => Vec::new(),
-        };
-
-        state_with_config(SamplingConfig {
-            rules,
-            mode,
-            next_id: None,
-        })
-    }
-
     fn samplingresult_from_rules_and_proccessing_flag(
         rules: Vec<SamplingRule>,
         processing_enabled: bool,
@@ -740,6 +645,54 @@ mod tests {
         assert_eq!(spec.unwrap().unwrap().sample_rate, 0.2);
     }
 
+    #[test]
+    fn test_sample_rate() {
+        let event_state_drop = state_with_rule_and_condition(
+            Some(0.0),
+            RuleType::Transaction,
+            SamplingMode::Received,
+            RuleCondition::all(),
+        );
+
+        let envelope = new_envelope(true, "foo");
+
+        let event = Event {
+            id: Annotated::new(EventId::new()),
+            ty: Annotated::new(EventType::Transaction),
+            transaction: Annotated::new("foo".to_owned()),
+            ..Event::default()
+        };
+
+        // if it matches the transaction rule, the transaction should be dropped
+        let should_drop = should_keep_event(
+            envelope.dsc(),
+            Some(&event),
+            None,
+            &event_state_drop,
+            Some(&event_state_drop),
+            false,
+        );
+
+        assert!(matches!(should_drop, SamplingResult::Drop(_)));
+
+        let event_state_keep = state_with_rule_and_condition(
+            Some(1.0),
+            RuleType::Transaction,
+            SamplingMode::Received,
+            RuleCondition::all(),
+        );
+
+        let should_keep = should_keep_event(
+            envelope.dsc(),
+            Some(&event),
+            None,
+            &event_state_keep,
+            Some(&event_state_keep),
+            false,
+        );
+
+        assert!(matches!(should_keep, SamplingResult::Keep));
+    }
     #[test]
     fn test_event_decaying_rule_with_linear_function() {
         let now = Utc::now();