Fix flaky test with proper synchronization using internal method

Author: steingran

src/MultiLock/LeaderElectionService.cs

@@ -680,8 +680,7 @@ private async Task BroadcastEventsAsync()
     {
         // Read until the channel is completed (not until cancelled)
         // This ensures all pending events are broadcast during graceful shutdown
-        // Note: ReadAllAsync() completes gracefully when the channel writer is completed,
-        // so no try-catch is needed for normal operation
+        // Note: ReadAllAsync() completes gracefully when the channel writer is completed
         await foreach (LeadershipChangedEventArgs eventArgs in broadcastChannel.Reader.ReadAllAsync().ConfigureAwait(false))
         {
             // Broadcast to all subscriber channels
@@ -705,8 +704,17 @@ private async Task BroadcastEventsAsync()
     /// An async enumerable that yields leadership change events as they occur.
     /// The enumeration will continue until the service is disposed or the cancellation token is triggered.
     /// </returns>
-    public async IAsyncEnumerable<LeadershipChangedEventArgs> GetLeadershipChangesAsync(
-        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    public IAsyncEnumerable<LeadershipChangedEventArgs> GetLeadershipChangesAsync(
+        CancellationToken cancellationToken = default)
+        => GetLeadershipChangesAsyncCore(cancellationToken, subscriberRegistered: null);
+
+    /// <summary>
+    /// Internal overload that signals when the subscriber is registered.
+    /// Used by tests to avoid race conditions.
+    /// </summary>
+    internal async IAsyncEnumerable<LeadershipChangedEventArgs> GetLeadershipChangesAsyncCore(
+        [EnumeratorCancellation] CancellationToken cancellationToken,
+        TaskCompletionSource? subscriberRegistered)
     {
         ThrowIfDisposed();
 
@@ -723,6 +731,9 @@ public async IAsyncEnumerable<LeadershipChangedEventArgs> GetLeadershipChangesAs
             subscriberChannels.Add(subscriberChannel);
         }
 
+        // Signal that the subscriber is now registered
+        subscriberRegistered?.TrySetResult();
+
         try
         {
             // Read from the subscriber channel

src/MultiLock/MultiLock.csproj

@@ -13,4 +13,8 @@
     <PackageReference Include="Microsoft.Extensions.Options" />
   </ItemGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="MultiLock.Tests" />
+  </ItemGroup>
+
 </Project>

tests/MultiLock.Tests/AsyncEnumerableApiTests.cs

@@ -324,16 +324,17 @@ public async Task GetLeadershipChangesAsync_CancellationToken_ShouldStopEnumerat
     {
         // Arrange
         ServiceProvider services = TestHelpers.CreateLeaderElectionService("test-participant");
-        ILeaderElectionService service = services.GetRequiredService<ILeaderElectionService>();
+        var service = (LeaderElectionService)services.GetRequiredService<ILeaderElectionService>();
         var events = new List<LeadershipChangedEventArgs>();
         object eventsLock = new();
         var cts = new CancellationTokenSource();
+        var subscriberRegistered = new TaskCompletionSource();
 
         // Act - Start listening BEFORE starting service
         CancellationToken cancellationToken = cts.Token;
         var eventTask = Task.Run(async () =>
         {
-            await foreach (LeadershipChangedEventArgs e in service.GetLeadershipChangesAsync(cancellationToken))
+            await foreach (LeadershipChangedEventArgs e in service.GetLeadershipChangesAsyncCore(cancellationToken, subscriberRegistered))
             {
                 lock (eventsLock)
                 {
@@ -342,8 +343,8 @@ public async Task GetLeadershipChangesAsync_CancellationToken_ShouldStopEnumerat
             }
         }, cancellationToken);
 
-        // Give the event listener task a chance to start
-        await Task.Delay(TimeSpan.FromMilliseconds(50), cts.Token);
+        // Wait for the subscriber to be registered (no race condition)
+        await subscriberRegistered.Task;
 
         await service.StartAsync(cts.Token);
         await service.WaitForLeadershipAsync(cts.Token);