test : Add test to confirm Context expected behavior (#2724)

opentelemetry-sdk/src/trace/mod.rs

@@ -76,6 +76,53 @@ mod tests {
         Context, KeyValue,
     };
 
+    #[test]
+    fn span_modification_via_context() {
+        let exporter = InMemorySpanExporterBuilder::new().build();
+        let provider = SdkTracerProvider::builder()
+            .with_span_processor(SimpleSpanProcessor::new(exporter.clone()))
+            .build();
+        let tracer = provider.tracer("test_tracer");
+
+        #[derive(Debug, PartialEq)]
+        struct ValueA(u64);
+
+        let span = tracer.start("span-name");
+
+        // start with Current, which should have no span
+        let cx = Context::current();
+        assert!(!cx.has_active_span());
+
+        // add span to context
+        let cx_with_span = cx.with_span(span);
+        assert!(cx_with_span.has_active_span());
+        assert!(!cx.has_active_span());
+
+        // modify the span by using span_ref from the context
+        // this is the only way to modify the span as span
+        // is moved to context.
+        let span_ref = cx_with_span.span();
+        span_ref.set_attribute(KeyValue::new("attribute1", "value1"));
+
+        // create a new context, which should not affect the original
+        let cx_with_span_and_more = cx_with_span.with_value(ValueA(1));
+
+        // modify the span again using the new context.
+        // this should still be using the original span itself.
+        let span_ref_new = cx_with_span_and_more.span();
+        span_ref_new.set_attribute(KeyValue::new("attribute2", "value2"));
+
+        span_ref_new.end();
+
+        let exported_spans = exporter
+            .get_finished_spans()
+            .expect("Spans are expected to be exported.");
+        // There should be a single span, with attributes from both modifications.
+        assert_eq!(exported_spans.len(), 1);
+        let span = &exported_spans[0];
+        assert_eq!(span.attributes.len(), 2);
+    }
+
     #[test]
     fn tracer_in_span() {
         // Arrange

opentelemetry/src/context.rs

@@ -385,6 +385,43 @@ impl Hasher for IdHasher {
 mod tests {
     use super::*;
 
+    #[test]
+    fn context_immutable() {
+        #[derive(Debug, PartialEq)]
+        struct ValueA(u64);
+        #[derive(Debug, PartialEq)]
+        struct ValueB(u64);
+
+        // start with Current, which should be an empty context
+        let cx = Context::current();
+        assert_eq!(cx.get::<ValueA>(), None);
+        assert_eq!(cx.get::<ValueB>(), None);
+
+        // with_value should return a new context,
+        // leaving the original context unchanged
+        let cx_new = cx.with_value(ValueA(1));
+
+        // cx should be unchanged
+        assert_eq!(cx.get::<ValueA>(), None);
+        assert_eq!(cx.get::<ValueB>(), None);
+
+        // cx_new should contain the new value
+        assert_eq!(cx_new.get::<ValueA>(), Some(&ValueA(1)));
+
+        // cx_new should be unchanged
+        let cx_newer = cx_new.with_value(ValueB(1));
+
+        // Cx and cx_new are unchanged
+        assert_eq!(cx.get::<ValueA>(), None);
+        assert_eq!(cx.get::<ValueB>(), None);
+        assert_eq!(cx_new.get::<ValueA>(), Some(&ValueA(1)));
+        assert_eq!(cx_new.get::<ValueB>(), None);
+
+        // cx_newer should contain both values
+        assert_eq!(cx_newer.get::<ValueA>(), Some(&ValueA(1)));
+        assert_eq!(cx_newer.get::<ValueB>(), Some(&ValueB(1)));
+    }
+
     #[test]
     fn nested_contexts() {
         #[derive(Debug, PartialEq)]