Function wait_for_message

tests/test_pubsub.py:32–52

(
    pubsub, timeout=0.5, ignore_subscribe_messages=False, node=None, func=None
)

def wait_for_message(
    pubsub, timeout=0.5, ignore_subscribe_messages=False, node=None, func=None
):
    now = time.monotonic()
    timeout = now + timeout
    while now < timeout:
        if node:
            message = pubsub.get_sharded_message(
                ignore_subscribe_messages=ignore_subscribe_messages, target_node=node
            )
        elif func:
            message = func(ignore_subscribe_messages=ignore_subscribe_messages)
        else:
            message = pubsub.get_message(
                ignore_subscribe_messages=ignore_subscribe_messages
            )
        if message is not None:
            return message
        time.sleep(0.01)
        now = time.monotonic()
    return None


def make_message(type, channel, data, pattern=None):