From 115d06e0a544f68b68fe31bf357dd01398c85baa Mon Sep 17 00:00:00 2001
From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com>
Date: Thu, 20 Mar 2025 15:45:21 +0000
Subject: [PATCH] Fix unstable telemetry UT

---
 .../KernelCore/KernelTests.cs                 | 22 +++++++++++++++----
 .../KernelCore/KernelTests.cs                 | 22 +++++++++++++++----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/dotnet/src/Connectors/Connectors.AzureOpenAI.UnitTests/KernelCore/KernelTests.cs b/dotnet/src/Connectors/Connectors.AzureOpenAI.UnitTests/KernelCore/KernelTests.cs
index 61685bb1daec..ae409be39872 100644
--- a/dotnet/src/Connectors/Connectors.AzureOpenAI.UnitTests/KernelCore/KernelTests.cs
+++ b/dotnet/src/Connectors/Connectors.AzureOpenAI.UnitTests/KernelCore/KernelTests.cs
@@ -77,6 +77,16 @@ public async Task FunctionUsageMetricsAreCapturedByTelemetryAsExpected()
             }
         });
 
+        var completed = false;
+
+        listener.MeasurementsCompleted = (instrument, state) =>
+        {
+            completed = true;
+            // Stop the listener to stop collecting data
+            Assert.Contains(12, measurements["semantic_kernel.function.invocation.token_usage.prompt"]);
+            Assert.Contains(5, measurements["semantic_kernel.function.invocation.token_usage.completion"]);
+        };
+
         listener.Start();  // Start the listener to begin collecting data
 
         this._multiMessageHandlerStub.ResponsesToReturn.Add(
@@ -90,12 +100,16 @@ public async Task FunctionUsageMetricsAreCapturedByTelemetryAsExpected()
 
         var kernelFunction = KernelFunctionFactory.CreateFromPrompt("prompt", loggerFactory: this._mockLoggerFactory.Object);
 
-        // Act
+        // Act & Assert
         var result = await kernel.InvokeAsync(kernelFunction);
 
-        // Assert
-        Assert.Contains(12, measurements["semantic_kernel.function.invocation.token_usage.prompt"]);
-        Assert.Contains(5, measurements["semantic_kernel.function.invocation.token_usage.completion"]);
+        listener.Dispose();
+
+        while (!completed)
+        {
+            // Wait for the measurements to be completed
+            await Task.Delay(100);
+        }
     }
 
     public void Dispose()
diff --git a/dotnet/src/Connectors/Connectors.OpenAI.UnitTests/KernelCore/KernelTests.cs b/dotnet/src/Connectors/Connectors.OpenAI.UnitTests/KernelCore/KernelTests.cs
index fdf17710b77c..4909504d743e 100644
--- a/dotnet/src/Connectors/Connectors.OpenAI.UnitTests/KernelCore/KernelTests.cs
+++ b/dotnet/src/Connectors/Connectors.OpenAI.UnitTests/KernelCore/KernelTests.cs
@@ -78,6 +78,16 @@ public async Task FunctionUsageMetricsAreCapturedByTelemetryAsExpected()
             }
         });
 
+        var completed = false;
+
+        listener.MeasurementsCompleted = (instrument, state) =>
+        {
+            completed = true;
+            // Stop the listener to stop collecting data
+            Assert.Contains(12, measurements["semantic_kernel.function.invocation.token_usage.prompt"]);
+            Assert.Contains(5, measurements["semantic_kernel.function.invocation.token_usage.completion"]);
+        };
+
         listener.Start();  // Start the listener to begin collecting data
 
         this._multiMessageHandlerStub.ResponsesToReturn.Add(
@@ -91,12 +101,16 @@ public async Task FunctionUsageMetricsAreCapturedByTelemetryAsExpected()
 
         var kernelFunction = KernelFunctionFactory.CreateFromPrompt("prompt", loggerFactory: this._mockLoggerFactory.Object);
 
-        // Act
+        // Act & Assert
         var result = await kernel.InvokeAsync(kernelFunction);
 
-        // Assert
-        Assert.Contains(12, measurements["semantic_kernel.function.invocation.token_usage.prompt"]);
-        Assert.Contains(5, measurements["semantic_kernel.function.invocation.token_usage.completion"]);
+        listener.Dispose();
+
+        while (!completed)
+        {
+            // Wait for the measurements to be completed
+            await Task.Delay(100);
+        }
     }
 
     public void Dispose()