From 35f08e46218eb09ca9a77622aa4d940ee55aa7e1 Mon Sep 17 00:00:00 2001
From: Peter Giacomo Lombardo <peter.lombardo@hivemq.com>
Date: Fri, 17 May 2024 10:24:06 +0200
Subject: [PATCH] Tests: Better TaskCompletion Strategy (#158)

---
 .../HiveMQClient/SubscribeBuilderTest.cs      | 27 ++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/Tests/HiveMQtt.Test/HiveMQClient/SubscribeBuilderTest.cs b/Tests/HiveMQtt.Test/HiveMQClient/SubscribeBuilderTest.cs
index d8e33057..0d680759 100644
--- a/Tests/HiveMQtt.Test/HiveMQClient/SubscribeBuilderTest.cs
+++ b/Tests/HiveMQtt.Test/HiveMQClient/SubscribeBuilderTest.cs
@@ -164,7 +164,13 @@ public async Task PerSubHandlerWithSingleLevelWildcardAsync()
         var connectResult = await subscribeClient.ConnectAsync().ConfigureAwait(false);
         Assert.True(connectResult.ReasonCode == ConnAckReasonCode.Success);
 
-        var tcs = new TaskCompletionSource<bool>();
+        // This is the fix for:
+        // per-subscription OnMessageReceivedEventLauncher exception: One or more errors occurred. (An attempt was
+        // made to transition a task to a final state when it had already completed.)
+        var tcs1 = new TaskCompletionSource<bool>();
+        var tcs2 = new TaskCompletionSource<bool>();
+        var tcs3 = new TaskCompletionSource<bool>();
+
         var messageCount = 0;
 
         var subscribeOptions = new SubscribeOptionsBuilder()
@@ -179,7 +185,18 @@ public async Task PerSubHandlerWithSingleLevelWildcardAsync()
 
                 if (messageCount == 3)
                 {
-                    tcs.SetResult(true);
+                    if (args.PublishMessage.Topic == "tests/PerSubHandlerWithSingleLevelWildcard/0/msg")
+                    {
+                        tcs1.SetResult(true);
+                    }
+                    else if (args.PublishMessage.Topic == "tests/PerSubHandlerWithSingleLevelWildcard/1/msg")
+                    {
+                        tcs2.SetResult(true);
+                    }
+                    else if (args.PublishMessage.Topic == "tests/PerSubHandlerWithSingleLevelWildcard/2/msg")
+                    {
+                        tcs3.SetResult(true);
+                    }
                 }
             })
             .Build();
@@ -200,11 +217,13 @@ public async Task PerSubHandlerWithSingleLevelWildcardAsync()
         }
 
         // Wait for the 3 messages to be received by the per-subscription handler
-        var handlerResult = await tcs.Task.WaitAsync(TimeSpan.FromSeconds(10)).ConfigureAwait(false);
-        Assert.True(handlerResult);
+        await Task.WhenAll(new Task[] { tcs1.Task, tcs2.Task, tcs3.Task }).ConfigureAwait(false);
 
         var disconnectResult = await subscribeClient.DisconnectAsync().ConfigureAwait(false);
         Assert.True(disconnectResult);
+
+        disconnectResult = await pubClient.DisconnectAsync().ConfigureAwait(false);
+        Assert.True(disconnectResult);
     }
 
     [Fact]