Fixes #28.

Only invoke start callback after the connection request. This way the connection info of the request is set when the callback is invoked.
statianzo · Jan 3, 2012 · 8c726e0 · 8c726e0
1 parent 25c371f
commit 8c726e0
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 10 deletions.
diff --git a/src/Fleck.Tests/WebSocketConnectionTests.cs b/src/Fleck.Tests/WebSocketConnectionTests.cs
@@ -17,8 +17,9 @@ public void Setup()
             _socketMock = new Mock<ISocket>();
             _handlerMock = new Mock<IHandler>();
             _connection = new WebSocketConnection(_socketMock.Object,
+                                                  connection => { },
                                                   b => new WebSocketHttpRequest(),
-                                                  r => _handlerMock.Object); 
+                                                  r => _handlerMock.Object);
         }
 
         [Test]
@@ -59,6 +60,53 @@ public void ShouldNotReadWhenSocketClosed()
             _socketMock.Verify(x => x.Receive(It.IsAny<byte[]>(), It.IsAny<Action<int>>(), It.IsAny<Action<Exception>>(), 0), Times.Never());
         }
 
+        [Test]
+        public void ShouldRaiseInitializeOnFirstRead()
+        {
+            bool initializeRaised = false;
+            var connection = new WebSocketConnection(_socketMock.Object,
+                                                  conn => { initializeRaised = true; },
+                                                  b => new WebSocketHttpRequest(),
+                                                  r => _handlerMock.Object);
+
+            _socketMock.SetupGet(x => x.Connected).Returns(true);
+            SetupReadLengths(1, 0);
+            connection.StartReceiving();
+
+            Assert.IsTrue(initializeRaised);
+        }
+
+        [Test]
+        public void ShouldNotRaiseInitializeIfParseRequestReturnsNull()
+        {
+            bool initializeRaised = false;
+            var connection = new WebSocketConnection(_socketMock.Object,
+                                                  conn => { initializeRaised = true; },
+                                                  b => null,
+                                                  r => _handlerMock.Object);
+
+            _socketMock.SetupGet(x => x.Connected).Returns(true);
+            SetupReadLengths(1, 0);
+            connection.StartReceiving();
+
+            Assert.IsFalse(initializeRaised);
+        }
+
+        [Test]
+        public void ShouldNotRaiseInitializeIfHandlerFactoryReturnsNull()
+        {
+            bool initializeRaised = false;
+            var connection = new WebSocketConnection(_socketMock.Object,
+                                                  conn => { initializeRaised = true; },
+                                                  b => new WebSocketHttpRequest(),
+                                                  r => null);
+
+            _socketMock.SetupGet(x => x.Connected).Returns(true);
+            SetupReadLengths(1, 0);
+            connection.StartReceiving();
+
+            Assert.IsFalse(initializeRaised);
+        }
 
         [Test]
         public void ShouldCallOnErrorWhenError()
@@ -79,7 +127,7 @@ public void ShouldCallOnErrorWhenError()
             _connection.StartReceiving();
             Assert.IsTrue(hit);
         }
-        
+
         [Test]
         public void ShouldSwallowObjectDisposedExceptionOnRead()
         {
@@ -99,7 +147,7 @@ public void ShouldSwallowObjectDisposedExceptionOnRead()
             _connection.StartReceiving();
             Assert.IsFalse(hit);
         }
-         
+
         private void SetupReadLengths(params int[] args)
         {
             var index = 0;
@@ -109,10 +157,10 @@ private void SetupReadLengths(params int[] args)
                 .Callback<byte[], Action<int>, Action<Exception>, int>((buffer, success, error, offset) =>
                 {
                     if (args.Length > index)
-                        success(args[index]);
+                        success(args[index++]);
                     else
                         _socketMock.SetupGet(x => x.Connected == false);
                 });
         }
     }
-}
+}
diff --git a/src/Fleck/WebSocketConnection.cs b/src/Fleck/WebSocketConnection.cs
@@ -6,19 +6,21 @@ namespace Fleck
 {
     public class WebSocketConnection : IWebSocketConnection
     {
-        public WebSocketConnection(ISocket socket, Func<byte[], WebSocketHttpRequest> parseRequest, Func<WebSocketHttpRequest, IHandler> handlerFactory)
+        public WebSocketConnection(ISocket socket, Action<IWebSocketConnection> initialize, Func<byte[], WebSocketHttpRequest> parseRequest, Func<WebSocketHttpRequest, IHandler> handlerFactory)
         {
             Socket = socket;
             OnOpen = () => { };
             OnClose = () => { };
             OnMessage = x => { };
             OnError = x => { };
+            _initialize = initialize;
             _handlerFactory = handlerFactory;
             _parseRequest = parseRequest;
         }
 
         public ISocket Socket { get; set; }
 
+        private readonly Action<IWebSocketConnection> _initialize;
         private readonly Func<WebSocketHttpRequest, IHandler> _handlerFactory;
         readonly Func<byte[], WebSocketHttpRequest> _parseRequest;
         public IHandler Handler { get; set; }
@@ -159,6 +161,8 @@ private void CreateHandler(IEnumerable<byte> data)
                 return;
             ConnectionInfo = WebSocketConnectionInfo.Create(request, Socket.RemoteIpAddress);
 
+            _initialize(this);
+
             var handshake = Handler.CreateHandshake();
             SendBytes(handshake, OnOpen);
         }
@@ -171,4 +175,4 @@ private void CloseSocket()
             Socket.Dispose();
         }
     }
-}
+}
diff --git a/src/Fleck/WebSocketServer.cs b/src/Fleck/WebSocketServer.cs
@@ -74,13 +74,11 @@ private void OnClientConnect(ISocket clientSocket)
 
             connection = new WebSocketConnection(
                 clientSocket,
+                _config,
                 bytes => RequestParser.Parse(bytes, _scheme),
                 r => HandlerFactory.BuildHandler(r, s => connection.OnMessage(s),
                                                          connection.Close));
 
-            _config(connection);
-
-
             if (IsSecure)
             {
                 FleckLog.Debug("Authenticating Secure Connection");