Add test token setter

Author: VianneyRuhlmann

data-pipeline/src/trace_exporter/mod.rs

@@ -757,6 +757,7 @@ pub struct TraceExporterBuilder {
     compute_stats_by_span_kind: bool,
     peer_tags: Vec<String>,
     telemetry: Option<TelemetryConfig>,
+    test_session_token: Option<String>,
 }
 
 impl TraceExporterBuilder {
@@ -873,6 +874,12 @@ impl TraceExporterBuilder {
         self
     }
 
+    /// Set the session token
+    pub fn set_test_session_token(&mut self, test_session_token: &str) -> &mut Self {
+        self.test_session_token = Some(test_session_token.to_string());
+        self
+    }
+
     /// Enable stats computation on traces sent through this exporter
     pub fn enable_stats(&mut self, bucket_size: Duration) -> &mut Self {
         self.stats_bucket_size = Some(bucket_size);
@@ -932,6 +939,7 @@ impl TraceExporterBuilder {
 
         let endpoint = Endpoint {
             url: agent_url,
+            test_token: self.test_session_token.map(|token| token.into()),
             ..Default::default()
         };
 

data-pipeline/tests/test_trace_exporter.rs

@@ -118,14 +118,15 @@ mod tracing_integration_tests {
         let task_result = task::spawn_blocking(move || {
             let mut builder = TraceExporter::builder();
             builder
-                .set_url(format!("{url}?test_session_token={snapshot_name}").as_str())
+                .set_url(url)
                 .set_language("test-lang")
                 .set_language_version("2.0")
                 .set_language_interpreter_vendor("vendor")
                 .set_language_interpreter("interpreter")
                 .set_tracer_version("1.0")
                 .set_env("test_env")
-                .set_service("test");
+                .set_service("test")
+                .set_test_session_token(snapshot_name);
 
             let trace_exporter = builder.build().expect("Unable to build TraceExporter");
 
@@ -165,14 +166,15 @@ mod tracing_integration_tests {
         let task_result = task::spawn_blocking(move || {
             let mut builder = TraceExporter::builder();
             builder
-                .set_url(format!("{url}?test_session_token={snapshot_name}").as_str())
+                .set_url(url)
                 .set_language("test-lang")
                 .set_language_version("2.0")
                 .set_language_interpreter_vendor("vendor")
                 .set_language_interpreter("interpreter")
                 .set_tracer_version("1.0")
                 .set_env("test_env")
                 .set_service("test")
+                .set_test_session_token(snapshot_name)
                 .set_input_format(TraceExporterInputFormat::V04)
                 .set_output_format(TraceExporterOutputFormat::V05);
             let trace_exporter = builder.build().expect("Unable to build TraceExporter");
@@ -206,14 +208,15 @@ mod tracing_integration_tests {
         let task_result = task::spawn_blocking(move || {
             let mut builder = TraceExporter::builder();
             builder
-                .set_url(format!("{url}?test_session_token={snapshot_name}").as_str())
+                .set_url(url)
                 .set_language("test-lang")
                 .set_language_version("2.0")
                 .set_language_interpreter_vendor("vendor")
                 .set_language_interpreter("interpreter")
                 .set_tracer_version("1.0")
                 .set_env("test_env")
                 .set_service("test")
+                .set_test_session_token(snapshot_name)
                 .set_input_format(TraceExporterInputFormat::V05)
                 .set_output_format(TraceExporterOutputFormat::V05);
             let trace_exporter = builder.build().expect("Unable to build TraceExporter");
@@ -274,13 +277,14 @@ mod tracing_integration_tests {
         let task_result = task::spawn_blocking(move || {
             let mut builder = TraceExporter::builder();
             builder
-                .set_url(format!("{url}?test_session_token={snapshot_name}").as_str())
+                .set_url(format!(url))
                 .set_language("test-lang")
                 .set_language_version("2.0")
                 .set_language_interpreter_vendor("vendor")
                 .set_language_interpreter("interpreter")
                 .set_tracer_version("1.0")
                 .set_env("test_env")
+                .set_test_session_token(snapshot_name)
                 .set_service("test");
 
             let trace_exporter = builder.build().expect("Unable to build TraceExporter");