This repository has been archived by the owner on Jul 15, 2023. It is now read-only.
/
socket_test.py
63 lines (52 loc) · 2.02 KB
/
socket_test.py
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
import asyncio
import datetime
import random
import websockets
class SocketTest(object):
def __init__(self, _event_loop, host='localhost', port=5678):
self.__connected = set()
ws_server = websockets.serve(self.web_sockets_handler, host, port)
_event_loop.run_until_complete(ws_server)
print('Host "%s" listening on port %i' % (host, port))
_event_loop.run_forever()
@asyncio.coroutine
def consumer(self, message):
print('Message: %s' % message)
@asyncio.coroutine
def producer(self):
now = datetime.datetime.utcnow().isoformat() + 'Z'
yield from asyncio.sleep(2)
return now
@asyncio.coroutine
def web_sockets_handler(self, websocket, path):
self.__connected.add(websocket)
print('Connected: %s' % websocket)
try:
yield from asyncio.wait([ws.send('client connected') for ws in self.__connected])
while True:
listener_task = asyncio.ensure_future(websocket.recv())
producer_task = asyncio.ensure_future(self.producer())
done, pending = yield from asyncio.wait(
[listener_task, producer_task],
return_when=asyncio.FIRST_COMPLETED)
if listener_task in done:
message = listener_task.result()
yield from self.consumer(message)
else:
listener_task.cancel()
if producer_task in done:
message = producer_task.result()
yield from websocket.send(message)
else:
producer_task.cancel()
except websockets.exceptions.ConnectionClosed:
print('Connection %s closed' % websocket)
finally:
self.__connected.remove(websocket)
if __name__ == '__main__':
config = {
'host': 'localhost',
'port': 5678
}
event_loop = asyncio.get_event_loop()
ws_test = SocketTest(event_loop, **config)