Add Test cases for EndpointUtilty, LifetimeManager and ServiceConnection

test/Microsoft.Azure.SignalR.Tests/Infrastructure/IHandshackMessageFactory.cs

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.Azure.SignalR.Protocol;
+
+namespace Microsoft.Azure.SignalR.Tests.Infrastructure
+{
+    interface IHandshackMessageFactory
+    {
+        ServiceMessage GetHandshackResposeMessage();
+    }
+}

test/Microsoft.Azure.SignalR.Tests/Infrastructure/PipeReaderExtensions.cs

@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
 using System;
 using System.Buffers;
 using System.Collections.Generic;

test/Microsoft.Azure.SignalR.Tests/Infrastructure/ServiceConnectionProxy.cs

@@ -9,6 +9,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.Azure.SignalR.Protocol;
+using Microsoft.Azure.SignalR.Tests.Infrastructure;
 using Microsoft.Extensions.Logging.Abstractions;
 using HandshakeRequestMessage = Microsoft.Azure.SignalR.Protocol.HandshakeRequestMessage;
 using HandshakeResponseMessage = Microsoft.Azure.SignalR.Protocol.HandshakeResponseMessage;
@@ -21,25 +22,36 @@ internal class ServiceConnectionProxy : IClientConnectionManager, IClientConnect
         private static readonly IServiceProtocol _serviceProtocol = new ServiceProtocol();
         private readonly PipeOptions _clientPipeOptions;
 
-        private IConnectionFactory ConnectionFactory { get; }
+        public IHandshackMessageFactory HandshackMessageFactory { get; }
+
+        public IConnectionFactory ConnectionFactory { get; }
 
         public IClientConnectionManager ClientConnectionManager { get; }
 
         public TestConnection ConnectionContext { get; }
 
-        private ServiceConnection ServiceConnection { get; }
+        public ServiceConnection ServiceConnection { get; }
 
         public ConcurrentDictionary<string, ServiceConnectionContext> ClientConnections => ClientConnectionManager.ClientConnections;
 
         private readonly ConcurrentDictionary<string, TaskCompletionSource<ConnectionContext>> _waitForConnectionOpen = new ConcurrentDictionary<string, TaskCompletionSource<ConnectionContext>>();
         private readonly ConcurrentDictionary<string, TaskCompletionSource<object>> _waitForConnectionClose = new ConcurrentDictionary<string, TaskCompletionSource<object>>();
+        private readonly ConcurrentDictionary<int, TaskCompletionSource<ConnectionContext>> _waitForServerConnection = new ConcurrentDictionary<int, TaskCompletionSource<ConnectionContext>>();
+        private readonly ConcurrentDictionary<Type, TaskCompletionSource<ServiceMessage>> _waitForMessage = new ConcurrentDictionary<Type, TaskCompletionSource<ServiceMessage>>();
+
+        private int _serverConnectionCount = 0;
 
-        public ServiceConnectionProxy(ConnectionDelegate callback = null, PipeOptions clientPipeOptions = null)
+        public ServiceConnectionProxy(ConnectionDelegate callback = null,
+            PipeOptions clientPipeOptions = null,
+            TestConnection connectionContext = null,
+            IConnectionFactory connectionFactory = null,
+            IHandshackMessageFactory handshackMessageFactory = null)
         {
-            ConnectionContext = new TestConnection();
-            ConnectionFactory = new TestConnectionFactory(ConnectionContext);
+            ConnectionContext = connectionContext ?? new TestConnection();
+            ConnectionFactory = connectionFactory ?? new TestConnectionFactory(ConnectionContext, this);
             ClientConnectionManager = new ClientConnectionManager();
             _clientPipeOptions = clientPipeOptions;
+            HandshackMessageFactory = handshackMessageFactory ?? new TestHandshackMessageFactory();
 
             ServiceConnection = new ServiceConnection(
                 _serviceProtocol,
@@ -53,8 +65,15 @@ public ServiceConnectionProxy(ConnectionDelegate callback = null, PipeOptions cl
 
         public Task StartAsync()
         {
-            _ = ServiceConnection.StartAsync();
-            return HandshakeAsync();
+            return ServiceConnection.StartAsync();
+        }
+
+        public Task ProcessIncomingAsync()
+        {
+            using (var processIncomingCts = new CancellationTokenSource(DefaultHandshakeTimeout))
+            {
+                return ProcessIncomingCoreAsync(ConnectionContext.Application.Input, processIncomingCts.Token);
+            }
         }
 
         public void Stop()
@@ -78,6 +97,17 @@ public Task WaitForConnectionCloseAsync(string connectionId)
             return _waitForConnectionClose.GetOrAdd(connectionId, key => new TaskCompletionSource<object>()).Task;
         }
 
+        public Task<ServiceMessage> WaitForMessageAsync(Type type)
+        {
+            return _waitForMessage.GetOrAdd(type, key => new TaskCompletionSource<ServiceMessage>()).Task;
+        }
+
+        public Task<ConnectionContext> WaitForServerConnectionAsync(int connectionCount)
+        {
+            return _waitForServerConnection.GetOrAdd(connectionCount, key => new TaskCompletionSource<ConnectionContext>())
+                .Task;
+        }
+
         private Task OnConnectionAsync(ConnectionContext connection)
         {
             var tcs = new TaskCompletionSource<object>();
@@ -112,14 +142,33 @@ public void RemoveClientConnection(string connectionId)
             }
         }
 
-        private async Task HandshakeAsync()
+        public void AddServerConnection()
+        {
+            Interlocked.Increment(ref _serverConnectionCount);
+
+            if (_waitForServerConnection.TryGetValue(_serverConnectionCount, out var tcs))
+            {
+                tcs.TrySetResult(null);
+            }
+        }
+
+        public void AddMessage(Type type, ServiceMessage message)
+        {
+            if (_waitForMessage.TryGetValue(type, out var tcs))
+            {
+                tcs.TrySetResult(message);
+            }
+        }
+
+        public async Task HandshakeAsync()
         {
             using (var handshakeCts = new CancellationTokenSource(DefaultHandshakeTimeout))
             {
                 await ReceiveHandshakeRequestAsync(ConnectionContext.Application.Input, handshakeCts.Token);
             }
 
-            await WriteMessageAsync(new HandshakeResponseMessage());
+            await WriteMessageAsync(HandshackMessageFactory.GetHandshackResposeMessage());
+            AddServerConnection();
         }
 
         private async Task ReceiveHandshakeRequestAsync(PipeReader input, CancellationToken cancellationToken)
@@ -170,6 +219,41 @@ private async Task ReceiveHandshakeRequestAsync(PipeReader input, CancellationTo
             }
         }
 
+        private async Task ProcessIncomingCoreAsync(PipeReader input, CancellationToken cancellationToken)
+        {
+            while (true)
+            {
+                var result = await input.ReadAsync(cancellationToken);
+                var buffer = result.Buffer;
+
+                var consumed = buffer.Start;
+                var examined = buffer.End;
+
+                try
+                {
+                    if (!buffer.IsEmpty)
+                    {
+                        if (_serviceProtocol.TryParseMessage(ref buffer, out var message))
+                        {
+                            consumed = buffer.Start;
+                            examined = consumed;
+
+                            AddMessage(message.GetType(), message);
+                        }
+                    }
+
+                    if (result.IsCompleted)
+                    {
+                        break;
+                    }
+                }
+                finally
+                {
+                    input.AdvanceTo(consumed, examined);
+                }
+            }
+        }
+
         public ServiceConnectionContext CreateConnection(OpenConnectionMessage message)
         {
             return new ServiceConnectionContext(message, _clientPipeOptions, _clientPipeOptions);

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestClientConnectionManager.cs

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Microsoft.Azure.SignalR.Tests.Infrastructure
+{
+    class TestClientConnectionManager : IClientConnectionManager
+    {
+        public void AddClientConnection(ServiceConnectionContext clientConnection)
+        {
+        }
+
+        public void RemoveClientConnection(string connectionId)
+        {
+        }
+
+        public ConcurrentDictionary<string, ServiceConnectionContext> ClientConnections { get; }
+    }
+}

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestConnection.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO.Pipelines;
@@ -9,7 +10,7 @@
 
 namespace Microsoft.Azure.SignalR.Tests
 {
-    public class TestConnection : ConnectionContext
+    public class TestConnection : ConnectionContext, IDisposable
     {
         public TestConnection()
         {
@@ -31,5 +32,19 @@ public TestConnection()
         public override IDuplexPipe Transport { get; set; }
 
         public IDuplexPipe Application { get; set; }
+
+        // To simulate the reconnection after disposed.
+        public void Reconnect()
+        {
+            var pipeOptions = new PipeOptions();
+            var pair = DuplexPipe.CreateConnectionPair(pipeOptions, pipeOptions);
+            Transport = pair.Transport;
+            Application = pair.Application;
+        }
+
+        public void Dispose()
+        {
+            Transport.Input.CancelPendingRead();
+        }
     }
 }

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestConnectionFactory.cs

@@ -7,17 +7,25 @@
 
 namespace Microsoft.Azure.SignalR.Tests
 {
-    public class TestConnectionFactory : IConnectionFactory
+    internal class TestConnectionFactory : IConnectionFactory
     {
         private readonly ConnectionContext _connection;
+        private readonly ServiceConnectionProxy _proxy;
 
-        public TestConnectionFactory(ConnectionContext connection)
+        public TestConnectionFactory(ConnectionContext connection, ServiceConnectionProxy proxy)
         {
             _connection = connection;
+            _proxy = proxy;
         }
 
         public Task<ConnectionContext> ConnectAsync(TransferFormat transferFormat, string connectionId, CancellationToken cancellationToken = default)
         {
+            if (_connection is TestConnection testConnection)
+            {
+                testConnection.Reconnect();
+            }
+            _ = _proxy.HandshakeAsync();
+
             return Task.FromResult(_connection);
         }
 

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestHandshackMessageFactory.cs

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Microsoft.Azure.SignalR.Protocol;
+
+namespace Microsoft.Azure.SignalR.Tests.Infrastructure
+{
+    class TestHandshackMessageFactory : IHandshackMessageFactory
+    {
+        private readonly string _errorMessage;
+
+        public TestHandshackMessageFactory(string errorMessage = null)
+        {
+            _errorMessage = errorMessage ?? string.Empty;
+        }
+
+        public ServiceMessage GetHandshackResposeMessage()
+        {
+            return new HandshakeResponseMessage(_errorMessage);
+        }
+    }
+}

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestHub.cs

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Microsoft.AspNetCore.SignalR;
+
+namespace Microsoft.Azure.SignalR.Tests.Infrastructure
+{
+    public class TestHub : Hub
+    {
+    }
+}

test/Microsoft.Azure.SignalR.Tests/Infrastructure/TestServiceConnectionManager.cs

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.SignalR;
+using Microsoft.Azure.SignalR.Protocol;
+
+namespace Microsoft.Azure.SignalR.Tests.Infrastructure
+{
+    class TestServiceConnectionManager<THub> : IServiceConnectionManager<THub> where THub : Hub
+    {
+        private readonly ConcurrentDictionary<Type, int> _writeAsyncCallCount = new ConcurrentDictionary<Type, int>();
+
+        public ServiceMessage ServiceMessage { get; private set; }
+
+        public void AddServiceConnection(ServiceConnection serviceConnection)
+        {
+        }
+
+        public async Task StartAsync()
+        {
+            await Task.CompletedTask;
+        }
+
+        public async Task WriteAsync(ServiceMessage serviceMessage)
+        {
+            _writeAsyncCallCount.AddOrUpdate(serviceMessage.GetType(), 1, (_, value) => value + 1);
+            ServiceMessage = serviceMessage;
+            await Task.CompletedTask;
+        }
+
+        public int GetCallCount(Type type)
+        {
+            if (_writeAsyncCallCount.TryGetValue(type, out var count))
+            {
+                return count;
+            }
+
+            return 0;
+        }
+    }
+}