async def test_add_period(self):
AsyncExclusivePeriod.add_period(self, 'test4')
self.assertEqual(AsyncExclusivePeriod.get_obj_present_period(self),
'test1')
obj_period_names = AsyncExclusivePeriod.get_obj_period_names(self)
self.assertEqual({'test1', 'test2', 'test3', 'test4'},
set(obj_period_names))
async def _handle_wsq(self):
'''
始终使用原生ws连接队列中的最后一个
:return:
'''
_handle_raw_ws_msg_task = None
while not self._exiting:
# 拿到最新原生ws连接
while True:
self._wsq.previous_ws = self._wsq.present_ws
self._wsq.present_ws = await self._wsq.get()
self._wsq.task_done()
# 关闭之前的ws连接
if isinstance(self._wsq.previous_ws,
websockets.WebSocketClientProtocol):
asyncio.create_task(self._wsq.previous_ws.close())
# 拿完了，得到最后一个
if self._wsq.qsize() <= 0:
if _handle_raw_ws_msg_task:
# 删除之前的ws处理任务
asyncio.create_task(
ensureTaskCanceled(_handle_raw_ws_msg_task))
break
_handle_raw_ws_msg_task = asyncio.create_task(
self._handle_raw_ws_msg(self._wsq.present_ws))
self._handle_raw_ws_msg_task = _handle_raw_ws_msg_task
AsyncExclusivePeriod.set_obj_period(self, 'handing_ws')
async def _ws_manager(self):
# 启动ws连接队列的消息对接handlers处理任务
self._handle_wsq_task = asyncio.create_task(self._handle_wsq())
initail = True
while not self._exiting or initail:
try:
initail = False
new_ws = await self._create_ws()
# 更新连接
self._wsq.put_nowait(new_ws)
# wait until ws is exchanged.
await AsyncExclusivePeriod.wait_enter_period(
self, 'handing_ws')
logger.debug('New ws connection opened.')
# 等待需要更新连接的信号或者连接报错
self._when2create_new_ws_task = asyncio.create_task(
self._when2create_new_ws())
await asyncio.wait([self._when2create_new_ws_task])
await asyncio.create_task(asyncio.sleep(1))
# 进入更换ws时期
AsyncExclusivePeriod.set_obj_period(self, 'exchanging_ws')
if self._handle_raw_ws_msg_task:
self._handle_raw_ws_msg_task.cancel()
except asyncio.CancelledError:
raise
except:
await asyncio.create_task(asyncio.sleep(1))
async def wait_worker_accomplish_working_then_set_period_test1():
'''等待工人干完活就设置self到周期test1'''
while True:
wait_worker1_exit_working_task = asyncio.create_task(
AsyncExclusivePeriod.wait_exit_period(
'worker1', 'working'))
wait_worker2_exit_working_task = asyncio.create_task(
AsyncExclusivePeriod.wait_exit_period(
'worker2', 'working'))
await asyncio.gather(wait_worker1_exit_working_task,
wait_worker2_exit_working_task)
AsyncExclusivePeriod.set_obj_period(self, 'test1')
def __init__(self):
self._wsq: asyncio.Queue = asyncio.Queue()
self._wsq.present_ws: websockets.WebSocketClientProtocol = None
self._wsq.previous_ws: websockets.WebSocketClientProtocol = None
self._exiting = False
self._ws_manager_task = asyncio.create_task(self._ws_manager())
self._handlers = set()
# 更换ws连接时期和处理ws数据时期，刚开始是exchanging_ws时期
AsyncExclusivePeriod.create_obj_periods(self, 'exchanging_ws',
'handing_ws')
self._when2create_new_ws_task = None
self._handle_raw_ws_msg_task = None
self._handle_wsq_task: asyncio.Task = None
async def worker(name: str):
'''工人等self的test2时期，就工作，工作完再等'''
while True:
AsyncExclusivePeriod.set_obj_period(name, 'waiting')
await AsyncExclusivePeriod.wait_enter_period(self, 'test2')
AsyncExclusivePeriod.set_obj_period(name, 'working')
q.put_nowait({
"time": asyncio.get_running_loop().time(),
'msg': 'start to work'
})
await asyncio.sleep(1) # simulate work
q.put_nowait({
"time": asyncio.get_running_loop().time(),
'msg': 'end working'
}) # 测试间隔1s
async def _handle_raw_ws_msg(self, ws: websockets.WebSocketClientProtocol):
try:
async for msg in ws:
try:
# logger.debug(repr(old_ws_update_ws_task))
msg = await self._parse_raw_data(msg)
logger.debug('\n' + beeprint.pp(msg,
output=False,
string_break_enable=False,
sort_keys=False))
# logger.debug('\n' + repr(self._update_ws_task))
tasks = []
for handler in self._handlers:
if asyncio.iscoroutinefunction(handler):
tasks.append(
asyncio.create_task(handler(deepcopy(msg))))
else:
try:
handler(deepcopy(msg))
except:
pass
for task in tasks:
try:
await task
except:
pass
except:
logger.error('\n' + traceback.format_exc())
except asyncio.CancelledError:
# print('忽略CancelledError')
pass
except:
if self._when2create_new_ws_task and 'handing_ws' == AsyncExclusivePeriod.get_obj_present_period(
self):
self._when2create_new_ws_task.cancel()
logger.error(
'ERROR leads to update connection.\n' +
traceback.format_exc(), )
finally:
logger.debug('Old connection abandoned.')
async def test_in_test1_period_after_end_working():
await asyncio.sleep(0.5)
self.assertEqual(
'test1',
AsyncExclusivePeriod.get_obj_present_period(self))
async def setUp(self) -> None:
AsyncExclusivePeriod.create_obj_periods(self, 'test1', 'test2',
'test3')
self.assertEqual(AsyncExclusivePeriod.get_obj_present_period(self),
'test1')
async def set_test2():
'''只要离开test2时期就1s后设置到test2时期让工人干活'''
while True:
await AsyncExclusivePeriod.wait_outside_period(self, 'test2')
await asyncio.sleep(1)
AsyncExclusivePeriod.set_obj_period(self, 'test2') # 测试休息1s
async def test_handle_period_then_set_period(self):
q = asyncio.Queue()
# 两个工人
AsyncExclusivePeriod.create_obj_periods('worker1', 'waiting',
'working')
AsyncExclusivePeriod.create_obj_periods('worker2', 'waiting',
'working')
async def worker(name: str):
'''工人等self的test2时期，就工作，工作完再等'''
while True:
AsyncExclusivePeriod.set_obj_period(name, 'waiting')
await AsyncExclusivePeriod.wait_enter_period(self, 'test2')
AsyncExclusivePeriod.set_obj_period(name, 'working')
q.put_nowait({
"time": asyncio.get_running_loop().time(),
'msg': 'start to work'
})
await asyncio.sleep(1) # simulate work
q.put_nowait({
"time": asyncio.get_running_loop().time(),
'msg': 'end working'
}) # 测试间隔1s
# 测试在0.5s之后是test1时期
# 激活两工人
worker1_task = asyncio.create_task(worker('worker1'))
worker2_task = asyncio.create_task(worker('worker2'))
async def wait_worker_accomplish_working_then_set_period_test1():
'''等待工人干完活就设置self到周期test1'''
while True:
wait_worker1_exit_working_task = asyncio.create_task(
AsyncExclusivePeriod.wait_exit_period(
'worker1', 'working'))
wait_worker2_exit_working_task = asyncio.create_task(
AsyncExclusivePeriod.wait_exit_period(
'worker2', 'working'))
await asyncio.gather(wait_worker1_exit_working_task,
wait_worker2_exit_working_task)
AsyncExclusivePeriod.set_obj_period(self, 'test1')
wait_worker_accomplish_working_then_set_period_test1_task = \
asyncio.create_task(wait_worker_accomplish_working_then_set_period_test1())
async def noticing_in_test2():
'''提醒工人干活的时期'''
while True:
await AsyncExclusivePeriod.wait_inside_period(self, 'test2')
q.put_nowait({
"time": asyncio.get_running_loop().time(),
'msg': 'In test2, the workers should be working.'
})
await asyncio.sleep(0.05) # 测试是处于工作时期内
noticing_in_test2_task = asyncio.create_task(noticing_in_test2())
async def set_test2():
'''只要离开test2时期就1s后设置到test2时期让工人干活'''
while True:
await AsyncExclusivePeriod.wait_outside_period(self, 'test2')
await asyncio.sleep(1)
AsyncExclusivePeriod.set_obj_period(self, 'test2') # 测试休息1s
set_test2_task = asyncio.create_task(set_test2())
work_start_time = None
work_end_time = None
tasks = []
init_time = asyncio.get_running_loop().time()
while asyncio.get_running_loop().time() - init_time <= 5:
new_msg = await q.get()
if new_msg['msg'] == 'start to work':
work_start_time = new_msg['time']
if work_end_time:
self.assertEqual(round(work_start_time - work_end_time), 1)
elif new_msg['msg'] == 'end working':
work_end_time = new_msg['time']
self.assertEqual(round(work_end_time - work_start_time), 1)
async def test_in_test1_period_after_end_working():
await asyncio.sleep(0.5)
self.assertEqual(
'test1',
AsyncExclusivePeriod.get_obj_present_period(self))
tasks.append(
asyncio.create_task(
test_in_test1_period_after_end_working()))
elif new_msg['msg'] == 'In test2, the workers should be working.':
if work_end_time:
self.assertLessThan(new_msg['time'] - work_start_time, 1)
[await task for task in tasks]
async def _wait_then_set_period(self, obj, wait_seconds: float,
period_name: str):
await asyncio.sleep(wait_seconds)
AsyncExclusivePeriod.set_obj_period(obj, period_name)
async def test_set_get_obj_present_period(self):
AsyncExclusivePeriod.set_obj_period(self, 'test2')
self.assertEqual('test2',
AsyncExclusivePeriod.get_obj_present_period(self))
def test_get_obj_period_names(self):
obj_period_names = AsyncExclusivePeriod.get_obj_period_names(self)
self.assertEqual({'test1', 'test2', 'test3'}, set(obj_period_names))
def test_get_obj_period(self):
self.assertEqual(
AsyncExclusivePeriod._get_obj_period(self, 'test1')._name, 'test1')