Decouple the dispatcher from routing for easier testing

samples/SocketsSample/Startup.cs

@@ -31,11 +31,11 @@ namespace SocketsSample
                 app.UseDeveloperExceptionPage();
             }
 
-            app.UseSockets(d =>
+            app.UseSockets(routes =>
             {
-                d.MapSocketEndpoint<HubEndpoint>("/hubs");
-                d.MapSocketEndpoint<ChatEndPoint>("/chat");
-                d.MapSocketEndpoint<JsonRpcEndpoint>("/jsonrpc");
+                routes.MapSocketEndpoint<HubEndpoint>("/hubs");
+                routes.MapSocketEndpoint<ChatEndPoint>("/chat");
+                routes.MapSocketEndpoint<JsonRpcEndpoint>("/jsonrpc");
             });
         }
     }

src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs

@@ -2,10 +2,7 @@
 using System.Text;
 using System.Threading.Tasks;
 using Channels;
-using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Http;
-using Microsoft.AspNetCore.Routing;
-using Microsoft.AspNetCore.Sockets.Routing;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Primitives;
 
@@ -15,19 +12,6 @@ namespace Microsoft.AspNetCore.Sockets
     {
         private readonly ConnectionManager _manager = new ConnectionManager();
         private readonly ChannelFactory _channelFactory = new ChannelFactory();
-        private readonly RouteBuilder _routes;
-
-        public HttpConnectionDispatcher(IApplicationBuilder app)
-        {
-            _routes = new RouteBuilder(app);
-        }
-
-        public void MapSocketEndpoint<TEndPoint>(string path) where TEndPoint : EndPoint
-        {
-            _routes.AddPrefixRoute(path, new RouteHandler(c => Execute<TEndPoint>(path, c)));
-        }
-
-        public IRouter GetRouter() => _routes.Build();
 
         public async Task Execute<TEndPoint>(string path, HttpContext context) where TEndPoint : EndPoint
         {

src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs

@@ -1,19 +1,40 @@
 using System;
+using Microsoft.AspNetCore.Routing;
 using Microsoft.AspNetCore.Sockets;
+using Microsoft.AspNetCore.Sockets.Routing;
 
 namespace Microsoft.AspNetCore.Builder
 {
     public static class HttpDispatcherAppBuilderExtensions
     {
-        public static IApplicationBuilder UseSockets(this IApplicationBuilder app, Action<HttpConnectionDispatcher> callback)
+        public static IApplicationBuilder UseSockets(this IApplicationBuilder app, Action<SocketRouteBuilder> callback)
         {
-            var dispatcher = new HttpConnectionDispatcher(app);
-            callback(dispatcher);
+            var dispatcher = new HttpConnectionDispatcher();
+            var routes = new RouteBuilder(app);
+
+            callback(new SocketRouteBuilder(routes, dispatcher));
 
             // TODO: Use new low allocating websocket API
             app.UseWebSockets();
-            app.UseRouter(dispatcher.GetRouter());
+            app.UseRouter(routes.Build());
             return app;
         }
+
+        public class SocketRouteBuilder
+        {
+            private readonly HttpConnectionDispatcher _dispatcher;
+            private readonly RouteBuilder _routes;
+
+            public SocketRouteBuilder(RouteBuilder routes, HttpConnectionDispatcher dispatcher)
+            {
+                _routes = routes;
+                _dispatcher = dispatcher;
+            }
+
+            public void MapSocketEndpoint<TEndPoint>(string path) where TEndPoint : EndPoint
+            {
+                _routes.AddPrefixRoute(path, new RouteHandler(c => _dispatcher.Execute<TEndPoint>(path, c)));
+            }
+        }
     }
 }