2020-06-18 09:11:42 +03:00
|
|
|
import asyncio
|
2021-01-06 07:13:39 +03:00
|
|
|
import logging
|
2021-03-10 05:27:27 +03:00
|
|
|
import time
|
2020-12-01 11:37:22 +03:00
|
|
|
from typing import Callable
|
2020-06-18 09:11:42 +03:00
|
|
|
|
2021-04-04 06:55:26 +03:00
|
|
|
from chia.protocols.protocol_message_types import ProtocolMessageTypes
|
2021-02-09 10:14:37 +03:00
|
|
|
|
2021-01-06 07:13:39 +03:00
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
2020-06-18 09:35:13 +03:00
|
|
|
|
2020-12-01 11:37:22 +03:00
|
|
|
async def time_out_assert_custom_interval(timeout: int, interval, function, value=True, *args, **kwargs):
|
2020-06-18 09:11:42 +03:00
|
|
|
start = time.time()
|
|
|
|
while time.time() - start < timeout:
|
|
|
|
if asyncio.iscoroutinefunction(function):
|
|
|
|
f_res = await function(*args, **kwargs)
|
|
|
|
else:
|
|
|
|
f_res = function(*args, **kwargs)
|
|
|
|
if value == f_res:
|
2021-05-11 08:05:45 +03:00
|
|
|
return None
|
2020-06-18 09:11:42 +03:00
|
|
|
await asyncio.sleep(interval)
|
2021-05-20 02:41:38 +03:00
|
|
|
assert False, "Timed assertion timed out"
|
2020-06-18 09:11:42 +03:00
|
|
|
|
2020-06-18 09:35:13 +03:00
|
|
|
|
2020-12-01 11:37:22 +03:00
|
|
|
async def time_out_assert(timeout: int, function, value=True, *args, **kwargs):
|
|
|
|
await time_out_assert_custom_interval(timeout, 0.05, function, value, *args, *kwargs)
|
2020-06-18 09:35:13 +03:00
|
|
|
|
2020-06-18 09:11:42 +03:00
|
|
|
|
2020-12-01 11:37:22 +03:00
|
|
|
async def time_out_assert_not_none(timeout: int, function, *args, **kwargs):
|
2020-06-18 09:11:42 +03:00
|
|
|
start = time.time()
|
|
|
|
while time.time() - start < timeout:
|
|
|
|
if asyncio.iscoroutinefunction(function):
|
|
|
|
f_res = await function(*args, **kwargs)
|
|
|
|
else:
|
|
|
|
f_res = function(*args, **kwargs)
|
|
|
|
if f_res is not None:
|
2021-05-11 08:05:45 +03:00
|
|
|
return None
|
2020-06-18 09:11:42 +03:00
|
|
|
await asyncio.sleep(0.05)
|
2021-05-20 02:41:38 +03:00
|
|
|
assert False, "Timed assertion timed out"
|
2020-12-01 11:37:22 +03:00
|
|
|
|
|
|
|
|
|
|
|
def time_out_messages(incoming_queue: asyncio.Queue, msg_name: str, count: int = 1) -> Callable:
|
|
|
|
async def bool_f():
|
|
|
|
if incoming_queue.qsize() < count:
|
|
|
|
return False
|
|
|
|
for _ in range(count):
|
2021-02-17 04:12:26 +03:00
|
|
|
response = (await incoming_queue.get())[0].type
|
2021-02-09 10:14:37 +03:00
|
|
|
if ProtocolMessageTypes(response).name != msg_name:
|
2021-01-20 07:04:13 +03:00
|
|
|
# log.warning(f"time_out_message: found {response} instead of {msg_name}")
|
2020-12-01 11:37:22 +03:00
|
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
|
|
return bool_f
|