From e5e8d1bee332e4f025121a13ae382f728112ce67 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Tue, 4 Oct 2016 00:12:07 -0700
Subject: [PATCH] Decouple the dispatcher from routing for easier testing

---
 samples/SocketsSample/Startup.cs              |  8 ++---
 .../HttpConnectionDispatcher.cs               | 16 ----------
 .../HttpDispatcherAppBuilderExtensions.cs     | 29 ++++++++++++++++---
 3 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/samples/SocketsSample/Startup.cs b/samples/SocketsSample/Startup.cs
index bf64b38ded..224379c195 100644
--- a/samples/SocketsSample/Startup.cs
+++ b/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");
             });
         }
     }
diff --git a/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs b/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs
index f5aea8747c..af9ff72a65 100644
--- a/src/Microsoft.AspNetCore.Sockets/HttpConnectionDispatcher.cs
+++ b/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
         {
diff --git a/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs b/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs
index f8165beec1..51204bd5e8 100644
--- a/src/Microsoft.AspNetCore.Sockets/HttpDispatcherAppBuilderExtensions.cs
+++ b/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)));
+            }
+        }
     }
 }