forked from MicroPythonOS/MicroPythonOS
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_websocket.py
More file actions
113 lines (92 loc) · 4.26 KB
/
test_websocket.py
File metadata and controls
113 lines (92 loc) · 4.26 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import asyncio
import unittest
import _thread
import time
from mpos import App, AppManager
from mpos import TaskManager
from websocket import WebSocketApp
class TestMutlipleWebsocketsAsyncio(unittest.TestCase):
max_allowed_connections = 3 # max that echo.websocket.org allows
#relays = ["wss://echo.websocket.org" ]
#relays = ["wss://echo.websocket.org", "wss://echo.websocket.org"]
#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}")
self.on_message_called = True
def on_open(self, wsapp):
print(f"on_open called: {wsapp}")
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!")
self.on_ping_called = True
def on_close(self, wsapp, close_status_code, close_msg):
print(f"on_close called: {wsapp}")
self.on_close_called += 1
def on_error(self, wsapp, arg1):
print(f"on_error called: {wsapp}, {arg1}")
self.on_error_called += 1
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 run_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(f"self.on_open_called: {self.on_open_called} so 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(f"self.on_close_called: {self.on_close_called} so 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.assertGreaterEqual(self.on_close_called, min(len(self.relays),self.max_allowed_connections), "on_close was called for less than allowed connections")
self.assertEqual(self.on_error_called, max(0, len(self.relays) - self.max_allowed_connections), "expecting one error per failed connection")
# Wait for *all* of them to finish (or be cancelled)
# If this hangs, it's also a failure:
print(f"doing gather of tasks: {tasks}")
for index, task in enumerate(tasks): print(f"task {index}: ph_key:{task.ph_key} done:{task.done()} running {task.coro}")
await asyncio.gather(*tasks, return_exceptions=True)
def wait_for_ping(self):
self.on_ping_called = False
for _ in range(60):
print("Waiting for on_ping to be called...")
if self.on_ping_called:
print("yes, it was called!")
break
time.sleep(1)
self.assertTrue(self.on_ping_called)
def test_it(self):
asyncio.run(self.run_main())