diff --git a/src/HotChocolate/Core/test/Subscriptions.Postgres.Tests/ResilientNpgsqlConnectionTests.cs b/src/HotChocolate/Core/test/Subscriptions.Postgres.Tests/ResilientNpgsqlConnectionTests.cs index bf39dc12ed3..3ee5069b1ad 100644 --- a/src/HotChocolate/Core/test/Subscriptions.Postgres.Tests/ResilientNpgsqlConnectionTests.cs +++ b/src/HotChocolate/Core/test/Subscriptions.Postgres.Tests/ResilientNpgsqlConnectionTests.cs @@ -101,7 +101,19 @@ public async Task DisposeAsync_Should_Disconnect() public async Task Reconnect_Should_ReconnectWhenConnectionIsClosed() { // Arrange - Func onConnect = _ => ValueTask.CompletedTask; + var onConnectCalled = 0; + var reconnected = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + ValueTask onConnect(CancellationToken _) + { + // Initialize connects once; the reconnect after the drop is the second connect. + if (++onConnectCalled == 2) + { + reconnected.TrySetResult(); + } + + return ValueTask.CompletedTask; + } Func onDisconnect = _ => ValueTask.CompletedTask; var resilientNpgsqlConnection = @@ -114,9 +126,8 @@ public async Task Reconnect_Should_ReconnectWhenConnectionIsClosed() await connection!.CloseAsync(); // Assert - SpinWait.SpinUntil( - () => connection != resilientNpgsqlConnection.Connection, - TimeSpan.FromSeconds(1)); + await reconnected.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.NotEqual(connection, resilientNpgsqlConnection.Connection); }