tests: create activity on consuming message

Farfetch · Nov 14, 2023 · 3e926e7 · joelfoliveira · Nov 14, 2023 · 3e926e7
1 parent d6d3f7e
commit 3e926e7
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 5 deletions.
diff --git a/src/KafkaFlow.IntegrationTests/Core/Middlewares/GzipMiddleware.cs b/src/KafkaFlow.IntegrationTests/Core/Middlewares/GzipMiddleware.cs
@@ -1,12 +1,15 @@
 namespace KafkaFlow.IntegrationTests.Core.Middlewares
 {
+    using System.Diagnostics;
     using System.Threading.Tasks;
     using KafkaFlow.IntegrationTests.Core.Handlers;
 
     internal class GzipMiddleware : IMessageMiddleware
     {
         public async Task Invoke(IMessageContext context, MiddlewareDelegate next)
         {
+            using var activity = ActivitySourceAccessor.ActivitySource.StartActivity("integration-test", ActivityKind.Internal);
+
             MessageStorage.Add((byte[]) context.Message.Value);
             await next(context);
         }

diff --git a/src/KafkaFlow.IntegrationTests/OpenTelemetryTests.cs b/src/KafkaFlow.IntegrationTests/OpenTelemetryTests.cs
@@ -56,14 +56,42 @@ public async Task AddOpenTelemetry_ProducingAndConsumingOneMessage_TraceAndSpans
             await producer.ProduceAsync(null, message);
 
             // Assert
-            var (producerSpan, consumerSpan) = await this.WaitForSpansAsync();
+            var (producerSpan, consumerSpan, internalSpan) = await this.WaitForSpansAsync();
 
             Assert.IsNotNull(this.exportedItems);
             Assert.IsNull(producerSpan.ParentId);
             Assert.AreEqual(producerSpan.TraceId, consumerSpan.TraceId);
             Assert.AreEqual(consumerSpan.ParentSpanId, producerSpan.SpanId);
         }
 
+        [TestMethod]
+        public async Task AddOpenTelemetry_CreateActivityOnConsumingMessage_TraceIsPropagatedToCreatedActivity()
+        {
+            // Arrange
+            var provider = await this.GetServiceProvider();
+            MessageStorage.Clear();
+
+            using var tracerProvider = Sdk.CreateTracerProviderBuilder()
+            .AddSource("KafkaFlow.OpenTelemetry")
+            .AddInMemoryExporter(this.exportedItems)
+            .Build();
+
+            var producer = provider.GetRequiredService<IMessageProducer<GzipProducer>>();
+            var message = this.fixture.Create<byte[]>();
+
+            // Act
+            await producer.ProduceAsync(null, message);
+
+            // Assert
+            var (producerSpan, consumerSpan, internalSpan) = await this.WaitForSpansAsync();
+
+            Assert.IsNotNull(this.exportedItems);
+            Assert.IsNull(producerSpan.ParentId);
+            Assert.AreEqual(producerSpan.TraceId, consumerSpan.TraceId);
+            Assert.AreEqual(consumerSpan.ParentSpanId, producerSpan.SpanId);
+            Assert.AreEqual(internalSpan.TraceId, consumerSpan.TraceId);
+        }
+
         [TestMethod]
         public async Task AddOpenTelemetry_ProducingAndConsumingOneMessage_BaggageIsPropagatedFromTestActivityToConsumer()
         {
@@ -97,7 +125,7 @@ public async Task AddOpenTelemetry_ProducingAndConsumingOneMessage_BaggageIsProp
             await producer.ProduceAsync(null, message);
 
             // Assert
-            var (producerSpan, consumerSpan) = await this.WaitForSpansAsync();
+            var (producerSpan, consumerSpan, internalSpan) = await this.WaitForSpansAsync();
 
             Assert.IsNotNull(this.exportedItems);
             Assert.AreEqual(producerSpan.TraceId, consumerSpan.TraceId);
@@ -181,9 +209,9 @@ await Policy
                 .ExecuteAsync(() => Task.FromResult(this.isPartitionAssigned));
         }
 
-        private async Task<(Activity producerSpan, Activity consumerSpan)> WaitForSpansAsync()
+        private async Task<(Activity producerSpan, Activity consumerSpan, Activity internalSpan)> WaitForSpansAsync()
         {
-            Activity producerSpan = null, consumerSpan = null;
+            Activity producerSpan = null, consumerSpan = null, internalSpan = null;
 
             await Policy
                 .HandleResult<bool>(isAvailable => !isAvailable)
@@ -192,11 +220,12 @@ await Policy
                 {
                     producerSpan = this.exportedItems.Find(x => x.Kind == ActivityKind.Producer);
                     consumerSpan = this.exportedItems.Find(x => x.Kind == ActivityKind.Consumer);
+                    internalSpan = this.exportedItems.Find(x => x.Kind == ActivityKind.Internal);
 
                     return Task.FromResult(producerSpan != null && consumerSpan != null);
                 });
 
-            return (producerSpan, consumerSpan);
+            return (producerSpan, consumerSpan, internalSpan);
         }
     }
 }