Refactored the three async modes as separate modules for greater flex…

…ibility

(Idea derived from pull request #1 by @drdaeman)

engineio/async_eventlet.py

@@ -0,0 +1,9 @@
+import importlib
+
+async = {
+    'threading': importlib.import_module('eventlet.green.threading'),
+    'queue': importlib.import_module('eventlet.queue'),
+    'queue_class': 'Queue',
+    'websocket': importlib.import_module('eventlet.websocket'),
+    'websocket_class': 'WebSocketWSGI'
+}

engineio/async_gevent.py

@@ -0,0 +1,9 @@
+import importlib
+
+async = {
+    'threading': importlib.import_module('gevent.threading'),
+    'queue': importlib.import_module('gevent.queue'),
+    'queue_class': 'Queue',
+    'websocket': None,
+    'websocket_class': None
+}

engineio/async_threading.py

@@ -0,0 +1,14 @@
+import importlib
+
+try:
+    queue = importlib.import_module('queue')
+except ImportError:  # pragma: no cover
+    queue = importlib.import_module('Queue')  # pragma: no cover
+
+async = {
+    'threading': importlib.import_module('threading'),
+    'queue': queue,
+    'queue_class': 'Queue',
+    'websocket': None,
+    'websocket_class': None
+}

engineio/server.py

@@ -73,44 +73,22 @@ def __init__(self, async_mode=None, ping_timeout=60, ping_interval=25,
             else:
                 self.logger.setLevel(logging.ERROR)
 
-        threading = None
-        queue = None
-        queue_class = None
-        websocket = None
-        if async_mode is None or async_mode == 'eventlet':
-            try:
-                threading = importlib.import_module('eventlet.green.threading')
-                queue = importlib.import_module('eventlet.queue')
-                queue_class = 'Queue'
-                websocket = importlib.import_module('eventlet.websocket')
-                async_mode = 'eventlet'
-            except ImportError:
-                pass
-        if async_mode is None or async_mode == 'gevent':
+        if async_mode is None:
+            modes = ['eventlet', 'gevent', 'threading']
+        else:
+            modes = [async_mode]
+        self.async = None
+        self.async_mode = None
+        for mode in modes:
             try:
-                threading = importlib.import_module('gevent.threading')
-                queue = importlib.import_module('gevent.queue')
-                queue_class = 'JoinableQueue'
-                websocket = None
-                async_mode = 'gevent'
+                self.async = importlib.import_module(
+                    'engineio.async_' + mode).async
+                self.async_mode = mode
+                break
             except ImportError:
                 pass
-        if async_mode is None or async_mode == 'threading':
-            threading = importlib.import_module('threading')
-            try:
-                queue = importlib.import_module('queue')
-            except ImportError:  # pragma: no cover
-                queue = importlib.import_module('Queue')  # pragma: no cover
-            queue_class = 'Queue'
-            websocket = None
-            async_mode = 'threading'
-        if threading is None:
+        if self.async_mode is None:
             raise ValueError('Invalid async_mode specified')
-        self.async_mode = async_mode
-        self.async = {'threading': threading,
-                      'queue': queue,
-                      'queue_class': queue_class,
-                      'websocket': websocket}
         self.logger.info('Server initialized for %s.', self.async_mode)
 
     def on(self, event, handler=None):
@@ -277,7 +255,7 @@ def _handle_connect(self, environ):
     def _upgrades(self, sid):
         """Return the list of possible upgrades for a client connection."""
         if not self.allow_upgrades or self._get_socket(sid).upgraded or \
-                self.async['websocket'] is None:
+                self.async['websocket_class'] is None:
             return []
         return ['websocket']
 

engineio/socket.py

@@ -101,8 +101,9 @@ def _upgrade_websocket(self, environ, start_response):
         """Upgrade the connection from polling to websocket."""
         if self.upgraded:
             raise IOError('Socket has been upgraded already')
-        ws = self.server.async['websocket'].WebSocketWSGI(
-            self._websocket_handler)
+        websocket_class = getattr(self.server.async['websocket'],
+                                  self.server.async['websocket_class'])
+        ws = websocket_class(self._websocket_handler)
         return ws(environ, start_response)
 
     def _websocket_handler(self, ws):

tests/test_server.py

@@ -1,4 +1,5 @@
 import gzip
+import importlib
 import logging
 import unittest
 import zlib
@@ -14,7 +15,23 @@
 from engineio import server
 
 
+original_import_module = importlib.import_module
+
+
+def _mock_import(module, pkg=None):
+    if module.startswith('engineio.'):
+        return original_import_module(module, pkg)
+    return module
+
+
 class TestServer(unittest.TestCase):
+    _mock_async = mock.MagicMock()
+    _mock_async.async = {
+        'threading': 't',
+        'queue': 'q',
+        'queue_class': 'qc',
+        'websocket_class': 'wc'}
+
     def _get_mock_socket(self):
         mock_socket = mock.MagicMock()
         mock_socket.closed = False
@@ -36,58 +53,67 @@ def test_create(self):
         for arg in six.iterkeys(kwargs):
             self.assertEqual(getattr(s, arg), kwargs[arg])
 
-    @mock.patch('importlib.import_module', side_effect=lambda mod: mod)
-    def test_async_mode_threading(self, import_module):
+    def test_async_mode_threading(self):
         s = server.Server(async_mode='threading')
         self.assertEqual(s.async_mode, 'threading')
-        self.assertEqual(s.async['threading'], 'threading')
-        self.assertEqual(s.async['queue'], 'queue')
+
+        import threading
+        try:
+            import queue
+        except ImportError:
+            import Queue as queue
+
+        self.assertEqual(s.async['threading'], threading)
+        self.assertEqual(s.async['queue'], queue)
+        self.assertEqual(s.async['queue_class'], 'Queue')
         self.assertEqual(s.async['websocket'], None)
+        self.assertEqual(s.async['websocket_class'], None)
 
-    @mock.patch('importlib.import_module', side_effect=lambda mod: mod)
-    def test_async_mode_eventlet(self, import_module):
+    def test_async_mode_eventlet(self):
         s = server.Server(async_mode='eventlet')
         self.assertEqual(s.async_mode, 'eventlet')
-        self.assertEqual(s.async['threading'], 'eventlet.green.threading')
-        self.assertEqual(s.async['queue'], 'eventlet.queue')
-        self.assertEqual(s.async['websocket'], 'eventlet.websocket')
 
-    @mock.patch('importlib.import_module', side_effect=lambda mod: mod)
+        from eventlet.green import threading
+        from eventlet import queue
+        from eventlet import websocket
+
+        self.assertEqual(s.async['threading'], threading)
+        self.assertEqual(s.async['queue'], queue)
+        self.assertEqual(s.async['queue_class'], 'Queue')
+        self.assertEqual(s.async['websocket'], websocket)
+        self.assertEqual(s.async['websocket_class'], 'WebSocketWSGI')
+
+    @mock.patch('importlib.import_module', side_effect=_mock_import)
     def test_async_mode_gevent(self, import_module):
         s = server.Server(async_mode='gevent')
         self.assertEqual(s.async_mode, 'gevent')
         self.assertEqual(s.async['threading'], 'gevent.threading')
         self.assertEqual(s.async['queue'], 'gevent.queue')
+        self.assertEqual(s.async['queue_class'], 'Queue')
         self.assertEqual(s.async['websocket'], None)
+        self.assertEqual(s.async['websocket_class'], None)
 
-    @mock.patch('importlib.import_module', side_effect=lambda mod: mod)
+    @mock.patch('importlib.import_module', side_effect=[ImportError])
     def test_async_mode_invalid(self, import_module):
         self.assertRaises(ValueError, server.Server, async_mode='foo')
 
-    @mock.patch('importlib.import_module', side_effect=lambda mod: mod)
+    @mock.patch('importlib.import_module', side_effect=[_mock_async])
     def test_async_mode_auto_eventlet(self, import_module):
         s = server.Server()
         self.assertEqual(s.async_mode, 'eventlet')
-        self.assertEqual(s.async['threading'], 'eventlet.green.threading')
-        self.assertEqual(s.async['queue'], 'eventlet.queue')
-        self.assertEqual(s.async['websocket'], 'eventlet.websocket')
 
-    @mock.patch('importlib.import_module', side_effect=[ImportError, 'a', 'b'])
+    @mock.patch('importlib.import_module', side_effect=[ImportError,
+                                                        _mock_async])
     def test_async_mode_auto_gevent(self, import_module):
         s = server.Server()
         self.assertEqual(s.async_mode, 'gevent')
-        self.assertEqual(s.async['threading'], 'a')
-        self.assertEqual(s.async['queue'], 'b')
-        self.assertEqual(s.async['websocket'], None)
 
     @mock.patch('importlib.import_module', side_effect=[ImportError,
-                                                        ImportError, 'a', 'b'])
+                                                        ImportError,
+                                                        _mock_async])
     def test_async_mode_auto_threading(self, import_module):
         s = server.Server()
         self.assertEqual(s.async_mode, 'threading')
-        self.assertEqual(s.async['threading'], 'a')
-        self.assertEqual(s.async['queue'], 'b')
-        self.assertEqual(s.async['websocket'], None)
 
     def test_generate_id(self):
         s = server.Server()

tests/test_socket.py

@@ -136,14 +136,15 @@ def test_upgrade_handshake(self):
     def test_upgrade(self):
         mock_server = self._get_mock_server()
         mock_server.async['websocket'] = mock.MagicMock()
+        mock_server.async['websocket_class'] = 'WebSocket'
         mock_ws = mock.MagicMock()
-        mock_server.async['websocket'].WebSocketWSGI.configure_mock(
+        mock_server.async['websocket'].WebSocket.configure_mock(
             return_value=mock_ws)
         s = socket.Socket(mock_server, 'sid')
         environ = "foo"
         start_response = "bar"
         s._upgrade_websocket(environ, start_response)
-        mock_server.async['websocket'].WebSocketWSGI.assert_called_once_with(
+        mock_server.async['websocket'].WebSocket.assert_called_once_with(
             s._websocket_handler)
         mock_ws.assert_called_once_with(environ, start_response)