added flask_test_client option to Socket.IO test client

flask_socketio/__init__.py

@@ -606,10 +606,28 @@ def sleep(self, seconds=0):
         """
         return self.server.sleep(seconds)
 
-    def test_client(self, app, namespace=None, query_string=None, headers=None):
-        """Return a simple SocketIO client that can be used for unit tests."""
+    def test_client(self, app, namespace=None, query_string=None,
+                    headers=None, flask_test_client=None):
+        """The Socket.IO test client is useful for testing a Flask-SocketIO
+        server. It works in a similar way to the Flask Test Client, but
+        adapted to the Socket.IO server.
+
+        :param app: The Flask application instance.
+        :param namespace: The namespace for the client. If not provided, the
+                          client connects to the server on the global
+                          namespace.
+        :param query_string: A string with custom query string arguments.
+        :param headers: A dictionary with custom HTTP headers.
+        :param flask_test_client: The instance of the Flask test client
+                                  currently in use. Passing the Flask test
+                                  client is optional, but is necessary if you
+                                  want the Flask user session and any other
+                                  cookies set in HTTP routes accessible from
+                                  Socket.IO events.
+        """
         return SocketIOTestClient(app, self, namespace=namespace,
-                                  query_string=query_string, headers=headers)
+                                  query_string=query_string, headers=headers,
+                                  flask_test_client=flask_test_client)
 
     def _handle_event(self, handler, message, namespace, sid, *args):
         if sid not in self.server.environ:

flask_socketio/test_client.py

@@ -16,12 +16,18 @@ class SocketIOTestClient(object):
                       connects to the server on the global namespace.
     :param query_string: A string with custom query string arguments.
     :param headers: A dictionary with custom HTTP headers.
+    :param flask_test_client: The instance of the Flask test client
+                              currently in use. Passing the Flask test
+                              client is optional, but is necessary if you
+                              want the Flask user session and any other
+                              cookies set in HTTP routes accessible from
+                              Socket.IO events.
     """
     queue = {}
     acks = {}
 
     def __init__(self, app, socketio, namespace=None, query_string=None,
-                 headers=None):
+                 headers=None, flask_test_client=None):
         def _mock_send_packet(sid, pkt):
             if pkt.packet_type == packet.EVENT or \
                     pkt.packet_type == packet.BINARY_EVENT:
@@ -41,6 +47,7 @@ def _mock_send_packet(sid, pkt):
                                   'namespace': pkt.namespace or '/'}
 
         self.app = app
+        self.flask_test_client = flask_test_client
         self.sid = uuid.uuid4().hex
         self.queue[self.sid] = []
         self.acks[self.sid] = None
@@ -77,6 +84,9 @@ def connect(self, namespace=None, query_string=None, headers=None):
             url += query_string
         environ = EnvironBuilder(url, headers=headers).get_environ()
         environ['flask.app'] = self.app
+        if self.flask_test_client:
+            # inject cookies from Flask
+            self.flask_test_client.cookie_jar.inject_wsgi(environ)
         self.socketio.server._handle_eio_connect(self.sid, environ)
         if namespace is not None and namespace != '/':
             pkt = packet.Packet(packet.CONNECT, namespace=namespace)

test_socketio.py

@@ -230,6 +230,12 @@ def on_other_custom_event(self, data):
 socketio.on_namespace(MyNamespace('/ns'))
 
 
+@app.route('/session')
+def session_route():
+    session['foo'] = 'bar'
+    return ''
+
+
 class TestSocketIO(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -411,14 +417,16 @@ def test_broadcast_namespace(self):
         self.assertEqual(len(client3.get_received()), 0)
 
     def test_session(self):
-        client = socketio.test_client(app)
+        flask_client = app.test_client()
+        flask_client.get('/session')
+        client = socketio.test_client(app, flask_test_client=flask_client)
         client.get_received()
         client.send('echo this message back')
         self.assertEqual(socketio.server.environ[client.sid]['saved_session'],
-                         {})
+                         {'foo': 'bar'})
         client.send('test session')
         self.assertEqual(socketio.server.environ[client.sid]['saved_session'],
-                         {'a': 'b'})
+                         {'a': 'b', 'foo': 'bar'})
 
     def test_room(self):
         client1 = socketio.test_client(app)