websocket: use asyncio instead of threads

tests/test_websocket.py

@@ -1,3 +1,4 @@
+import asyncio
 import unittest
 import _thread
 import time
@@ -7,14 +8,19 @@ import mpos.apps
 
 from websocket import WebSocketApp
 
-class TestWebsocket(unittest.TestCase):
+class TestMutlipleWebsocketsAsyncio(unittest.TestCase):
 
-    ws = None
+    max_allowed_connections = 3 # max that echo.websocket.org allows
 
-    on_open_called = None
-    on_message_called = None
-    on_ping_called = None
-    on_close_called = None
+    #relays = ["wss://echo.websocket.org", "wss://echo.websocket.org", "wss://echo.websocket.org" ] # more gives "too many requests" error
+    relays = ["wss://echo.websocket.org", "wss://echo.websocket.org", "wss://echo.websocket.org", "wss://echo.websocket.org", "wss://echo.websocket.org" ] # more might give "too many requests" error
+    wslist = []
+
+    on_open_called = 0
+    on_message_called = 0
+    on_ping_called = 0
+    on_close_called = 0
+    on_error_called = 0
 
     def on_message(self, wsapp, message: str):
         print(f"on_message received: {message}")
@@ -22,8 +28,8 @@ class TestWebsocket(unittest.TestCase):
         
     def on_open(self, wsapp):
         print(f"on_open called: {wsapp}")
-        self.on_open_called = True
-        self.ws.send('{"type": "subscribe","product_ids": ["BTC-USD"],"channels": ["ticker_batch"]}')
+        self.on_open_called += 1
+        #wsapp.send('{"type": "subscribe","product_ids": ["BTC-USD"],"channels": ["ticker_batch"]}')
 
     def on_ping(wsapp, message):
         print("Got a ping!")
@@ -31,19 +37,63 @@ class TestWebsocket(unittest.TestCase):
 
     def on_close(self, wsapp, close_status_code, close_msg):
         print(f"on_close called: {wsapp}")
-        self.on_close_called = True
+        self.on_close_called += 1
 
-    def websocket_thread(self):
-        wsurl = "wss://ws-feed.exchange.coinbase.com"
+    def on_error(self, wsapp, arg1):
+        print(f"on_error called: {wsapp}, {arg1}")
+        self.on_error_called += 1
 
-        self.ws = WebSocketApp(
-            wsurl,
-            on_open=self.on_open,
-            on_close=self.on_close,
-            on_message=self.on_message,
-            on_ping=self.on_ping
-        ) # maybe add other callbacks to reconnect when disconnected etc.
-        self.ws.run_forever()
+    async def closeall(self):
+        await asyncio.sleep(1)
+
+        self.on_close_called = 0
+        print("disconnecting...")
+        for ws in self.wslist:
+            await ws.close()
+
+    async def main(self) -> None:
+        tasks = []
+        self.wslist = []
+        for idx, wsurl in enumerate(self.relays):
+            print(f"creating WebSocketApp for {wsurl}")
+            ws = WebSocketApp(
+                wsurl,
+                on_open=self.on_open,
+                on_close=self.on_close,
+                on_message=self.on_message,
+                on_ping=self.on_ping,
+                on_error=self.on_error
+            )
+            print(f"creating task for {wsurl}")
+            tasks.append(asyncio.create_task(ws.run_forever(),))
+            print(f"created task for {wsurl}")
+            self.wslist.append(ws)
+
+        print(f"Starting {len(tasks)} concurrent WebSocket connections…")
+        await asyncio.sleep(2)
+        await self.closeall()
+
+        for _ in range(10):
+            print("Waiting for on_open to be called...")
+            if self.on_open_called == min(len(self.relays),self.max_allowed_connections):
+                print("yes, it was called!")
+                break
+            await asyncio.sleep(1)
+        self.assertTrue(self.on_open_called == min(len(self.relays),self.max_allowed_connections))
+
+        for _ in range(10):
+            print("Waiting for on_close to be called...")
+            if self.on_close_called == min(len(self.relays),self.max_allowed_connections):
+                print("yes, it was called!")
+                break
+            await asyncio.sleep(1)
+        self.assertTrue(self.on_close_called == min(len(self.relays),self.max_allowed_connections))
+
+        self.assertTrue(self.on_error_called == min(len(self.relays),self.max_allowed_connections))
+
+        # Wait for *all* of them to finish (or be cancelled)
+        # If this hangs, it's also a failure:
+        await asyncio.gather(*tasks, return_exceptions=True)
 
     def wait_for_ping(self):
         self.on_ping_called = False
@@ -55,39 +105,12 @@ class TestWebsocket(unittest.TestCase):
             time.sleep(1)
         self.assertTrue(self.on_ping_called)
 
-    def test_it(self):
-        on_open_called = False
-        _thread.stack_size(mpos.apps.good_stack_size())
-        _thread.start_new_thread(self.websocket_thread, ())
-        
-        self.on_open_called = False
-        self.on_message_called = False # message might be received very quickly, before we expect it
-        for _ in range(5):
-            print("Waiting for on_open to be called...")
-            if self.on_open_called:
-                print("yes, it was called!")
-                break
-            time.sleep(1)
-        self.assertTrue(self.on_open_called)
+    def test_it_loop(self):
+        for testnr in range(1):
+            print(f"starting iteration {testnr}")
+            asyncio.run(self.do_two())
+            print(f"finished iteration {testnr}")
 
-        self.on_message_called = False # message might be received very quickly, before we expect it
-        for _ in range(5):
-            print("Waiting for on_message to be called...")
-            if self.on_message_called:
-                print("yes, it was called!")
-                break
-            time.sleep(1)
-        self.assertTrue(self.on_message_called)
+    def do_two(self):
+        await self.main()
 
-        # Disabled because not all servers send pings:
-        # self.wait_for_ping()
-
-        self.on_close_called = False
-        self.ws.close()
-        for _ in range(5):
-            print("Waiting for on_close to be called...")
-            if self.on_close_called:
-                print("yes, it was called!")
-                break
-            time.sleep(1)
-        self.assertTrue(self.on_close_called)