fix: Correct flaky tests

Author: steingran

tests/MultiLock.Tests/InMemoryProviderTests.cs

@@ -77,7 +77,7 @@ public async Task GetCurrentLeaderAsync_WhenLeaderExists_ShouldReturnLeaderInfo(
         // Assert
         leader.ShouldNotBeNull();
         leader.LeaderId.ShouldBe(participantId);
-        leader.Metadata.ShouldContainKey("key");
+        leader.Metadata.ContainsKey("key").ShouldBeTrue();
         leader.Metadata["key"].ShouldBe("value");
     }
 
@@ -147,7 +147,7 @@ public async Task UpdateHeartbeatAsync_WhenIsLeader_ShouldSucceed()
 
         LeaderInfo? leader = await provider.GetCurrentLeaderAsync(electionGroup);
         leader.ShouldNotBeNull();
-        leader.Metadata.ShouldContainKey("key");
+        leader.Metadata.ContainsKey("key").ShouldBeTrue();
         leader.Metadata["key"].ShouldBe("updated-value");
     }
 

tests/MultiLock.Tests/LeadershipExtensionMethodsTests.cs

@@ -210,15 +210,14 @@ public async Task DistinctUntilChanged_ShouldFilterConsecutiveDuplicates()
         ILeaderElectionService service = services.GetRequiredService<ILeaderElectionService>();
         var events = new List<LeadershipChangedEventArgs>();
         object eventsLock = new();
-        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(10));
 
         // Act - Start listening BEFORE starting service
         CancellationToken cancellationToken = cts.Token;
         var eventTask = Task.Run(async () =>
         {
             await foreach (LeadershipChangedEventArgs e in service.GetLeadershipChangesAsync(cancellationToken)
-                .DistinctUntilChanged(cancellationToken)
-                .Take(2, cancellationToken))
+                .DistinctUntilChanged(cancellationToken))
             {
                 lock (eventsLock)
                 {
@@ -228,12 +227,18 @@ public async Task DistinctUntilChanged_ShouldFilterConsecutiveDuplicates()
         }, cancellationToken);
 
         // Give the async enumerable consumer time to start
-        await Task.Delay(TimeSpan.FromMilliseconds(50), cts.Token);
+        await Task.Delay(TimeSpan.FromMilliseconds(100), cancellationToken);
 
-        await service.StartAsync(cts.Token);
-        await service.WaitForLeadershipAsync(cts.Token);
-        await service.StopAsync(cts.Token);
-        await TestHelpers.WaitForConditionAsync(() => events.Count >= 1, TimeSpan.FromSeconds(5), cts.Token, eventsLock);
+        await service.StartAsync(cancellationToken);
+        await service.WaitForLeadershipAsync(cancellationToken);
+
+        // Wait for the first event (gained leadership)
+        await TestHelpers.WaitForConditionAsync(() => events.Count >= 1, TimeSpan.FromSeconds(3), cancellationToken, eventsLock);
+
+        await service.StopAsync(cancellationToken);
+
+        // Wait for the second event (lost leadership)
+        await TestHelpers.WaitForConditionAsync(() => events.Count >= 2, TimeSpan.FromSeconds(3), cancellationToken, eventsLock);
 
         // Cleanup - Cancel and wait for event task to complete before asserting
         await cts.CancelAsync();
@@ -245,7 +250,9 @@ public async Task DistinctUntilChanged_ShouldFilterConsecutiveDuplicates()
         {
             eventSnapshot = events.ToArray();
         }
-        eventSnapshot.Length.ShouldBeGreaterThanOrEqualTo(1);
+        eventSnapshot.Length.ShouldBeGreaterThanOrEqualTo(2);
+        eventSnapshot[0].BecameLeader.ShouldBeTrue();
+        eventSnapshot[1].LostLeadership.ShouldBeTrue();
 
         await services.DisposeAsync();
     }

tests/MultiLock.Tests/MultiInstanceLeaderElectionTests.cs

@@ -66,6 +66,9 @@ public async Task MultipleInstances_WithInMemoryProvider_ShouldElectOneLeader()
                 services.Add(service);
             }
 
+            // Give event subscriptions time to start
+            await Task.Delay(TimeSpan.FromMilliseconds(100), cts.Token);
+
             // Act - Start all instances
             Task[] startTasks = services.Select(s => s.StartAsync(cts.Token)).ToArray();
             await Task.WhenAll(startTasks);
@@ -76,6 +79,19 @@ await TestHelpers.WaitForConditionAsync(
                 TimeSpan.FromSeconds(5),
                 cts.Token);
 
+            // Wait for leadership events to be received
+            await TestHelpers.WaitForConditionAsync(
+                () =>
+                {
+                    lock (leadershipEvents)
+                    {
+                        return leadershipEvents.Any(e => e.IsLeader);
+                    }
+                },
+                TimeSpan.FromSeconds(3),
+                cts.Token,
+                leadershipEvents);
+
             // Assert - Exactly one leader should be elected
             var currentLeaders = services.Where(s => s.IsLeader).ToList();
             currentLeaders.Count.ShouldBe(1, "exactly one instance should be the leader");