diff --git a/datadog-trace-utils/src/send_data/mod.rs b/datadog-trace-utils/src/send_data/mod.rs index acc4fb2ff7..b09102cf53 100644 --- a/datadog-trace-utils/src/send_data/mod.rs +++ b/datadog-trace-utils/src/send_data/mod.rs @@ -78,6 +78,7 @@ pub enum Compression { None, } +#[derive(Clone)] pub struct SendDataBuilder { pub(crate) tracer_payloads: TracerPayloadCollection, pub(crate) size: usize, @@ -114,6 +115,16 @@ impl SendDataBuilder { self } + pub fn with_api_key(mut self, api_key: &str) -> SendDataBuilder { + self.target.api_key = Some(api_key.to_string().into()); + self + } + + pub fn with_retry_strategy(mut self, retry_strategy: RetryStrategy) -> SendDataBuilder { + self.retry_strategy = retry_strategy; + self + } + pub fn build(self) -> SendData { SendData { tracer_payloads: self.tracer_payloads, @@ -420,6 +431,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::send_with_retry::{RetryBackoffType, RetryStrategy}; use crate::test_utils::create_test_no_alloc_span; use crate::trace_utils::{construct_trace_chunk, construct_tracer_payload, RootSpanTags}; use crate::tracer_header_tags::TracerHeaderTags; @@ -1030,4 +1042,29 @@ mod tests { #[cfg(feature = "compression")] assert!(matches!(new_data.compression, Compression::None)); } + + #[test] + fn test_builder() { + let header_tags = HEADER_TAGS; + let payload = setup_payload(&header_tags); + let retry_strategy = RetryStrategy::new(5, 100, RetryBackoffType::Constant, None); + + let send_data = SendDataBuilder::new( + 100, + TracerPayloadCollection::V07(vec![payload]), + header_tags, + &Endpoint::default(), + ) + // Test with_api_key() + .with_api_key("TEST-KEY") + // Test with_retry_strategy() + .with_retry_strategy(retry_strategy.clone()) + .build(); + + assert_eq!( + send_data.target.api_key, + Some(std::borrow::Cow::Borrowed("TEST-KEY")) + ); + assert_eq!(send_data.retry_strategy, retry_strategy); + } }