Fix tests

test/test_coros.py

@@ -6,6 +6,35 @@
 from wsstat.clients import WebsocketTestingClient
 
 class TestCoros(object):
+    @asyncio.coroutine
+    def bad_websocket_handshake(*args, **kwargs):
+        raise InvalidHandshake("Bad status code: 200")
+
+    @asyncio.coroutine
+    def invalid_websocket_uri(*args, **kwargs):
+        raise InvalidURI("Invalid URI specified!")
+
+    @asyncio.coroutine
+    def mock_sleep(*args, **kwargs):
+        return
+
+    @asyncio.coroutine
+    def mocked_websocket(*args, **kwargs):
+        class MockedWebsocket(mock.Mock):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.count = 0
+
+            @asyncio.coroutine
+            def recv(self, *args, **kwargs):
+                self.count += 1
+                if self.count > 3:
+                    raise ConnectionClosed(1000, "Peace out homie!")
+                return "Hello World!"
+
+        return MockedWebsocket()
+
+
     def setup(self):
         self.identifier = "00000000"
 
@@ -18,64 +47,65 @@ def setup(self):
             )
         )
 
+    @mock.patch('wsstat.clients.websockets.connect', bad_websocket_handshake)
+    @mock.patch('wsstat.clients.asyncio.sleep', mock_sleep)
     def test_bad_handshake(self):
         event_loop = asyncio.new_event_loop()
         try:
             assert len(asyncio.Task.all_tasks(event_loop)) == 0
 
             # Test bad handshakes
-            with mock.patch('wsstat.clients.websockets.connect', bad_websocket_handshake):
-                event_loop.run_until_complete(self.client.create_websocket_connection())
+            event_loop.run_until_complete(self.client.create_websocket_connection())
 
-                # Make sure that the logger was called
-                assert self.client.logger.log.call_count > 0
-                # Ensure that there was a failed connection
-                assert isinstance(self.client.sockets[self.identifier], BaseException)
+            # Make sure that the logger was called
+            assert self.client.logger.log.call_count > 0
+            # Ensure that there was a failed connection
+            assert isinstance(self.client.sockets[self.identifier], BaseException)
         finally:
             event_loop.stop()
 
+    @mock.patch('wsstat.clients.websockets.connect', invalid_websocket_uri)
+    @mock.patch('wsstat.clients.asyncio.sleep', mock_sleep)
     def test_invalid_ws_url(self):
         event_loop = asyncio.new_event_loop()
         try:
             assert len(asyncio.Task.all_tasks(event_loop)) == 0
 
-            # Test bad handshakes
-            with mock.patch('wsstat.clients.websockets.connect', invalid_websocket_uri):
-                event_loop.run_until_complete(self.client.create_websocket_connection())
+            event_loop.run_until_complete(self.client.create_websocket_connection())
 
-                # Make sure that the logger was called
-                assert self.client.logger.log.call_count > 0
-                # Ensure that there was a failed connection
-                assert self.client.sockets[self.identifier] == False
+            # Make sure that the logger was called
+            assert self.client.logger.log.call_count > 0
+            # Ensure that there was a failed connection
+            assert self.client.sockets[self.identifier] == False
         finally:
             event_loop.stop()
 
+    @mock.patch('wsstat.clients.websockets.connect', mocked_websocket)
     def test_valid_connection(self):
         event_loop = asyncio.new_event_loop()
         try:
             # Test invalid URI
-            with mock.patch('wsstat.clients.websockets.connect', mocked_websocket):
-                event_loop.run_until_complete(self.client.create_websocket_connection())
-
-                # Ensure logging was successful
-                all_messages = self.arg_list_to_string(self.client.logger.log.call_args_list)
-                assert "Connecting" in all_messages
-                assert "Connected" in all_messages
-                assert "connection is closed" in all_messages
-                assert self.identifier in all_messages
-
-                connected_websocket = self.client.sockets[self.identifier]
-                assert connected_websocket.message_count == 3
-                assert connected_websocket.id == self.identifier
-                assert connected_websocket.last_message_recv != 0
-                assert connected_websocket.ws is not None
-
-                assert self.client.before_connect.call_count == 1
-                assert self.client.setup_websocket_connection.call_count == 1
-                assert self.client.get_identifier.call_count == 1
-                assert self.client.after_connect.call_count == 1
-                assert self.client.before_recv.call_count == 4
-                assert self.client.after_recv.call_count == 3
+            event_loop.run_until_complete(self.client.create_websocket_connection())
+
+            # Ensure logging was successful
+            all_messages = self.arg_list_to_string(self.client.logger.log.call_args_list)
+            assert "Connecting" in all_messages
+            assert "Connected" in all_messages
+            assert "connection is closed" in all_messages
+            assert self.identifier in all_messages
+
+            connected_websocket = self.client.sockets[self.identifier]
+            assert connected_websocket.message_count == 3
+            assert connected_websocket.id == self.identifier
+            assert connected_websocket.last_message_recv != 0
+            assert connected_websocket.ws is not None
+
+            assert self.client.before_connect.call_count == 1
+            assert self.client.setup_websocket_connection.call_count == 1
+            assert self.client.get_identifier.call_count == 1
+            assert self.client.after_connect.call_count == 1
+            assert self.client.before_recv.call_count == 4
+            assert self.client.after_recv.call_count == 3
         finally:
             event_loop.stop()
 
@@ -97,28 +127,4 @@ def _wrap_client(self, client):
         return client
 
     def arg_list_to_string(self, arglist):
-        return " | ".join([x[0][0] for x in arglist])
-
-@asyncio.coroutine
-def bad_websocket_handshake(*args, **kwargs):
-    raise InvalidHandshake("Bad status code: 200")
-
-@asyncio.coroutine
-def invalid_websocket_uri(*args, **kwargs):
-    raise InvalidURI("Invalid URI specified!")
-
-@asyncio.coroutine
-def mocked_websocket(*args, **kwargs):
-    class MockedWebsocket(mock.Mock):
-        def __init__(self, **kwargs):
-            super().__init__(**kwargs)
-            self.count = 0
-
-        @asyncio.coroutine
-        def recv(self, *args, **kwargs):
-            self.count += 1
-            if self.count > 3:
-                raise ConnectionClosed(1000, "Peace out homie!")
-            return "Hello World!"
-
-    return MockedWebsocket()
+        return " | ".join([x[0][0] for x in arglist])

wsstat/clients.py

@@ -17,6 +17,10 @@
 from wsstat.gui import BlinkBoardWidget, LoggerWidget
 
 import logging
+import sys
+
+if sys.version_info < (3, 4, 4):
+    asyncio.ensure_future = asyncio.async
 
 class ConnectedWebsocketConnection(object):
     def __init__(self, ws, identifier):
@@ -150,10 +154,7 @@ def create_websocket_connection(self):
 
                             return False
 
-
-                yield from asyncio.sleep(.25)
-
-
+                yield from asyncio.sleep(.25, loop=self.loop)
 
             # Create our handler object
             connected_websocket = ConnectedWebsocketConnection(websocket, identifier)

wsstat/demo.py

@@ -4,13 +4,18 @@
 import datetime
 import random
 import websockets
+import sys
 
-async def echo_time(websocket, path):
+if sys.version_info < (3, 4, 4):
+    asyncio.ensure_future = asyncio.async
+
+@asyncio.coroutine
+def echo_time(websocket, path):
     while True:
         now = datetime.datetime.utcnow().isoformat() + 'Z'
         try:
-            await websocket.send(now)
-            await asyncio.sleep(random.random() * 3)
+            yield from websocket.send(now)
+            yield from asyncio.sleep(random.random() * 3)
         except:
             pass