ConnectAsync should wait ConnectedMessage

Author: doghappy

src/SocketIOClient/V2/SocketIO.cs

@@ -58,17 +58,29 @@ private Uri ServerUri
 
     private readonly Dictionary<int, Action<IAckMessage>> _ackHandlers = new();
     private readonly Dictionary<int, Func<IAckMessage, Task>> _funcHandlers = new();
+    private TaskCompletionSource<Exception> _connCompletionSource = new();
     public SocketIOOptions Options { get; }
     public event EventHandler<Exception> OnReconnectError;
     public event EventHandler OnPing;
     public event EventHandler<TimeSpan> OnPong;
 
     public async Task ConnectAsync()
     {
-        await ConnectAsync(CancellationToken.None);
+        await ConnectAsync(CancellationToken.None).ConfigureAwait(false);
     }
 
     public async Task ConnectAsync(CancellationToken cancellationToken)
+    {
+        _ = ConnectCoreAsync(cancellationToken).ConfigureAwait(false);
+        var task = Task.Run(async () => await _connCompletionSource.Task.ConfigureAwait(false), cancellationToken);
+        var ex = await task.ConfigureAwait(false);
+        if (ex != null)
+        {
+            throw ex;
+        }
+    }
+
+    private async Task ConnectCoreAsync(CancellationToken cancellationToken)
     {
         var attempts = Options.Reconnection ? Options.ReconnectionAttempts : 1;
         for (int i = 0; i < attempts; i++)
@@ -81,9 +93,10 @@ public async Task ConnectAsync(CancellationToken cancellationToken)
             using var cts = new CancellationTokenSource(Options.ConnectionTimeout);
             try
             {
-                await session.ConnectAsync(cts.Token);
+                await session.ConnectAsync(cts.Token).ConfigureAwait(false);
                 _session = session;
                 _session.Subscribe(this);
+                // _sessionCompletionSource.SetResult(true);
             }
             catch (Exception e)
             {
@@ -92,10 +105,11 @@ public async Task ConnectAsync(CancellationToken cancellationToken)
                 OnReconnectError?.Invoke(this, ex);
                 if (i == attempts - 1)
                 {
+                    _connCompletionSource.SetResult(ex);
                     throw ex;
                 }
                 var delay = Random.Next(Options.ReconnectionDelayMax);
-                await Task.Delay(delay, CancellationToken.None);
+                await Task.Delay(delay, CancellationToken.None).ConfigureAwait(false);
             }
         }
     }
@@ -163,12 +177,12 @@ private async Task HandleAckMessage(IMessage message)
         }
     }
 
-    private Task HandleConnectedMessage(IMessage message)
+    private async Task HandleConnectedMessage(IMessage message)
     {
         var connectedMessage = (ConnectedMessage)message;
         Id = connectedMessage.Sid;
         Connected = true;
-        return Task.CompletedTask;
+        _connCompletionSource.SetResult(null);
     }
 
     private void HandlePongMessage(IMessage message)

tests/SocketIOClient.UnitTests/V2/SocketIOTests.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics;
 using FluentAssertions;
 using NSubstitute;
 using NSubstitute.ExceptionExtensions;
@@ -106,24 +107,26 @@ await _io
     [Fact]
     public async Task ConnectAsync_SessionSuccessfullyConnected_SessionSubscribeIO()
     {
-        await _io.ConnectAsync();
+        await ConnectAsync();
         _session.Received(1).Subscribe(_io);
     }
 
-    [Fact]
-    public async Task ConnectAsync_SessionSuccessfullyConnectedButNoConnectedMessage_ConnectedIsFalse()
+    private async Task ConnectAsync()
     {
-        await _io.ConnectAsync();
-        _io.Connected.Should().BeFalse();
+        await ConnectAsync(0);
     }
 
-    private async Task ConnectAsync()
+    private async Task ConnectAsync(int ms)
     {
-        await _io.ConnectAsync();
-        await _io.OnNextAsync(new ConnectedMessage
+        _ = Task.Run(async () =>
         {
-            Sid = "123",
+            await Task.Delay(ms);
+            await _io.OnNextAsync(new ConnectedMessage
+            {
+                Sid = "123",
+            });
         });
+        await _io.ConnectAsync();
     }
 
     [Fact]
@@ -134,6 +137,18 @@ public async Task ConnectAsync_ConnectedMessageReceived_ConnectedIsTrueIdHasValu
         _io.Id.Should().Be("123");
     }
 
+    [Fact]
+    public async Task ConnectAsync_ConnectedMessageDelay_ConnectAsyncIsSync()
+    {
+        var stopwatch = Stopwatch.StartNew();
+        await ConnectAsync(200);
+        stopwatch.Stop();
+
+        stopwatch.ElapsedMilliseconds.Should()
+            .BeGreaterThanOrEqualTo(200)
+            .And.BeLessThan(280);
+    }
+
     [Fact]
     public async Task ConnectAsyncCancellationToken_GivenACanceledToken_ThrowConnectionException()
     {