From fe41412eeb78e34f9cfd6dcb113e665e18bd664f Mon Sep 17 00:00:00 2001
From: John Luo <jt.luo@mail.utoronto.ca>
Date: Mon, 21 Dec 2015 16:37:49 -0800
Subject: [PATCH] Stardardizing option initialization using lambdas

---
 .../WebSocketMiddlewareExtensions.cs          | 22 +++++++++++++++++--
 test/AutobahnTestServer/Startup.cs            |  5 ++---
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/Microsoft.AspNet.WebSockets.Server/WebSocketMiddlewareExtensions.cs b/src/Microsoft.AspNet.WebSockets.Server/WebSocketMiddlewareExtensions.cs
index c676331771..b22a09ad7f 100644
--- a/src/Microsoft.AspNet.WebSockets.Server/WebSocketMiddlewareExtensions.cs
+++ b/src/Microsoft.AspNet.WebSockets.Server/WebSocketMiddlewareExtensions.cs
@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
 using Microsoft.AspNet.WebSockets.Server;
 
 namespace Microsoft.AspNet.Builder
@@ -9,11 +10,28 @@ namespace Microsoft.AspNet.Builder
     {
         public static IApplicationBuilder UseWebSockets(this IApplicationBuilder builder)
         {
-            return builder.UseWebSockets(new WebSocketOptions());
+            if (builder == null)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+
+            return builder.UseWebSockets(options => { });
         }
 
-        public static IApplicationBuilder UseWebSockets(this IApplicationBuilder builder, WebSocketOptions options) // TODO: [NotNull]
+        public static IApplicationBuilder UseWebSockets(this IApplicationBuilder builder, Action<WebSocketOptions> configureOptions)
         {
+            if (builder == null)
+            {
+                throw new ArgumentNullException(nameof(builder));
+            }
+            if (configureOptions == null)
+            {
+                throw new ArgumentNullException(nameof(configureOptions));
+            }
+
+            var options = new WebSocketOptions();
+            configureOptions(options);
+
             return builder.Use(next => new WebSocketMiddleware(next, options).Invoke);
         }
     }
diff --git a/test/AutobahnTestServer/Startup.cs b/test/AutobahnTestServer/Startup.cs
index 4ba3b887b7..22fd8aae83 100644
--- a/test/AutobahnTestServer/Startup.cs
+++ b/test/AutobahnTestServer/Startup.cs
@@ -8,7 +8,6 @@ using System.Threading.Tasks;
 using Microsoft.AspNet.Builder;
 using Microsoft.AspNet.Hosting;
 using Microsoft.AspNet.Http;
-using Microsoft.AspNet.WebSockets.Server;
 
 namespace AutobahnTestServer
 {
@@ -19,9 +18,9 @@ namespace AutobahnTestServer
             app.Map("/Managed", managedWebSocketsApp =>
             {
                 // Comment this out to test native server implementations
-                managedWebSocketsApp.UseWebSockets(new WebSocketOptions()
+                managedWebSocketsApp.UseWebSockets(options =>
                 {
-                    ReplaceFeature = true,
+                    options.ReplaceFeature = true;
                 });
 
                 managedWebSocketsApp.Use(async (context, next) =>