Python asyncio.wait_for() Examples

def _send_init_message_and_wait_ack(self) -> None:
        """Send init message to the provided websocket and wait for the connection ACK.

        If the answer is not a connection_ack message, we will return an Exception.
        """

        init_message = json.dumps(
            {"type": "connection_init", "payload": self.init_payload}
        )

        await self._send(init_message)

        # Wait for the connection_ack message or raise a TimeoutError
        init_answer = await asyncio.wait_for(self._receive(), self.ack_timeout)

        answer_type, answer_id, execution_result = self._parse_answer(init_answer)

        if answer_type != "connection_ack":
            raise TransportProtocolError(
                "Websocket server did not return a connection ack"
            ) 

def _clean_close(self, e: Exception) -> None:
        """Coroutine which will:

        - send stop messages for each active subscription to the server
        - send the connection terminate message
        """

        # Send 'stop' message for all current queries
        for query_id, listener in self.listeners.items():

            if listener.send_stop:
                await self._send_stop_message(query_id)
                listener.send_stop = False

        # Wait that there is no more listeners (we received 'complete' for all queries)
        try:
            await asyncio.wait_for(self._no_more_listeners.wait(), self.close_timeout)
        except asyncio.TimeoutError:  # pragma: no cover
            pass

        # Finally send the 'connection_terminate' message
        await self._send_connection_terminate_message() 

def test_election_early_wait_for_leadership(zk, path):
    elec = zk.recipes.LeaderElection(path)

    early_wait_success = asyncio.Event()

    async def wait_early():
        await elec.wait_for_leadership()
        assert elec.has_leadership
        early_wait_success.set()

    asyncio.create_task(wait_early())
    await asyncio.sleep(0.5)
    assert not elec.has_leadership

    await elec.volunteer()

    # NO WAIT
    await asyncio.wait_for(early_wait_success.wait(), timeout=0.5)

    await elec.resign()

    assert not elec.has_leadership

    await zk.delete(path) 

def test_reconnect_watcher(data_watcher, path, zk_disruptor, zk, zk2):
    test_data = uuid.uuid4().hex.encode()
    ready = data_watcher.client.loop.create_future()

    async def data_callback(d):
        print(f'Data callback get: {d}')
        if d == NoNode:
            return
        if d and not ready.done():
            print(f'Set result: {d} {ready}')
            ready.set_result(d)

    data_watcher.add_callback(path, data_callback)
    await zk_disruptor()
    await zk2.set_data(path, test_data)
    resp = await zk2.get_data(path)
    assert resp == test_data

    data = await asyncio.wait_for(ready, 1)
    assert data == test_data

    data_watcher.remove_callback(path, data_callback) 

def test_http_completion() -> None:
    # Ensure that the connecion callable returns on completion
    app = Quart(__name__)
    scope = {
        "headers": [(b"host", b"quart")],
        "http_version": "1.1",
        "method": "GET",
        "scheme": "https",
        "path": "/",
        "query_string": b"",
    }
    connection = ASGIHTTPConnection(app, scope)

    queue: asyncio.Queue = asyncio.Queue()
    queue.put_nowait({"type": "http.request", "body": b"", "more_body": False})

    async def receive() -> dict:
        # This will block after returning the first and only entry
        return await queue.get()

    async def send(message: dict) -> None:
        pass

    # This test fails if a timeout error is raised here
    await asyncio.wait_for(connection(receive, send), timeout=1) 

def test_data_watch_delete(zk, path, data_watcher):
    data = []
    ready = asyncio.Event()
    test_data = b'test'

    async def data_callback(d):
        data.append(d)
        ready.set()

    await zk.set_data(path, test_data)

    data_watcher.add_callback(path, data_callback)
    await asyncio.sleep(0.2)
    assert data == [test_data]
    ready.clear()
    await zk.delete(path)

    await asyncio.wait_for(ready.wait(), timeout=1)
    assert ready.is_set()
    assert data == [test_data, NoNode]
    data_watcher.remove_callback(path, data_callback)

    await zk.create(path) 

def test_data_watch(zk, path, data_watcher):
    data = []
    ready = asyncio.Event()
    test_data = b'test' * 1000

    async def data_callback(d):
        data.append(d)
        ready.set()

    data_watcher.add_callback(path, data_callback)
    assert data == []
    await zk.set_data(path, test_data)
    await asyncio.wait_for(ready.wait(), timeout=0.1)
    assert ready.is_set()
    assert data == [test_data]
    data_watcher.remove_callback(path, data_callback) 

def _parse_sdcard_list(self, done_cb):
        self.log.debug('Comms: _parse_sdcard_list')

        # setup callback to receive and parse listing data
        files = []
        f = asyncio.Future()
        self.redirect_incoming(lambda x: self._rcv_sdcard_line(x, files, f))

        # issue command
        self._write('M20\n')

        # wait for it to complete and get all the lines
        # add a long timeout in case it fails and we don't want to wait for ever
        try:
            await asyncio.wait_for(f, 10)

        except asyncio.TimeoutError:
            self.log.warning("Comms: Timeout waiting for sd card list")
            files = []

        self.redirect_incoming(None)

        # call upstream callback with results
        done_cb(files) 

def __connect(self):
        self.__check_closed()
        if self.connected:
            return
        try:
            self.logger.debug("Opening connection to %s:%d", self.host, self.port)
            future = asyncio.open_connection(self.host, self.port, loop=self.__loop)
            self.__reader, self.__writer = await asyncio.wait_for(
                future, timeout=self.connect_timeout, loop=self.__loop
            )

            await asyncio.wait_for(self.__connect_request_response(), timeout=self.request_timeout, loop=self.__loop)

            self.logger.debug("Socket connected successfully. Starting read loop.")
            self.connected = True
            self.__loop.create_task(self.__read_loop())
        except ConnectionError as e:
            self.logger.error("Connection error while connecting to server: %s", e)
            raise 

def _connect(self, timeout=None, ssl=None):
        await super()._connect(timeout=timeout, ssl=ssl)

        # Wait for EOF for 2 seconds (or if _wait_for_data's definition
        # is missing or different, just sleep for 2 seconds). This way
        # we give the proxy a chance to close the connection if the current
        # codec (which the proxy detects with the data we sent) cannot
        # be used for this proxy. This is a work around for #1134.
        # TODO Sleeping for N seconds may not be the best solution
        # TODO This fix could be welcome for HTTP proxies as well
        try:
            await asyncio.wait_for(self._reader._wait_for_data('proxy'), 2)
        except asyncio.TimeoutError:
            pass
        except Exception:
            await asyncio.sleep(2)

        if self._reader.at_eof():
            await self.disconnect()
            raise ConnectionError(
                'Proxy closed the connection after sending initial payload') 

def _get_result(self, future, start_time, timeout, pending, target_id):
        due = self._total_due
        if timeout is None:
            timeout = self._timeout

        if timeout is not None:
            due = min(due, start_time + timeout)

        # NOTE: We can't try/finally to pop from pending here because
        #       the event loop needs to get back to us, but it might
        #       dispatch another update before, and in that case a
        #       response could be set twice. So responses must be
        #       cleared when their futures are set to a result.
        return asyncio.wait_for(
            future,
            timeout=None if due == float('inf') else due - time.time(),
            loop=self._client.loop
        ) 

def close(self):
        if not self.started:
            log.debug('Do nothing because session is not started')
            return
        if self.closing:
            return
        self.closing = True
        if self.repair_loop_task:
            self.repair_loop_task.cancel()
            await asyncio.wait_for(self.send(protocol.CloseRequest()), self.timeout)
        if self.state.current_state != States.LOST:
            self.state.transition_to(States.LOST)
        if self.conn:
            await self.conn.close(self.timeout)
        self.closing = False
        self.started = False 

def heartbeat(self):
        if self.closing:
            return

        await self.ensure_safe_state()

        try:
            timeout = self.timeout - self.timeout/HEARTBEAT_FREQUENCY
            zxid, _ = await asyncio.wait_for(self.conn.send(protocol.PingRequest()), timeout)
            self.last_zxid = zxid
        except (exc.ConnectError, asyncio.TimeoutError):
            if self.state != States.SUSPENDED:
                self.state.transition_to(States.SUSPENDED)
        except Exception as e:
            log.exception('in heartbeat: {}'.format(e))
            raise e
        finally:
            self.set_heartbeat() 

def wait_on_sibling(self, sibling, timeout=None):
        deadline = Deadline(timeout)
        log.debug("Waiting on sibling %s", sibling)

        path = self.sibling_path(sibling)

        unblocked = self.client.wait_for_events([WatchEvent.DELETED], path)

        exists = await self.client.exists(path=path, watch=True)
        if not exists:
            unblocked.set_result(None)

        try:
            if not deadline.is_indefinite:
                await asyncio.wait_for(unblocked, deadline.timeout)
            else:
                await unblocked
        except asyncio.TimeoutError:
            raise exc.TimeoutError 

def delete_garbage_znodes(self, znode_label):
        MAXIMUM_WAIT = 60
        retry_policy = RetryPolicy.exponential_backoff(maximum=MAXIMUM_WAIT)
        while True:
            await self.client.session.state.wait_for(states.States.CONNECTED)
            await retry_policy.enforce()
            try:
                siblings = await self.get_siblings()
                for sibling in siblings:
                    if self.guid in sibling and self.determine_znode_label(
                            sibling) == znode_label:
                        path = self.sibling_path(sibling)
                        if path != self.owned_paths.get(znode_label, ''):
                            await self.client.delete(path)

                break
            except Exception:
                log.exception('Exception in delete_garbage_znodes:') 

def wait(self, timeout=None):
        deadline = Deadline(timeout)
        barrier_lifted = self.client.wait_for_events(
            [WatchEvent.DELETED], self.path
        )

        exists = await self.client.exists(path=self.path, watch=True)
        if not exists:
            return

        try:
            if not deadline.is_indefinite:
                await asyncio.wait_for(barrier_lifted, deadline.timeout)
            else:
                await barrier_lifted
        except asyncio.TimeoutError:
            raise exc.TimeoutError 

def wait_for_iterator(self, connection_observer, connection_observer_future):
        """
        Version of wait_for() intended to be used by Python3 to implement awaitable object.

        Note: we don't have timeout parameter here. If you want to await with timeout please do use asyncio machinery.
        For ex.:  await asyncio.wait_for(connection_observer, timeout=10)

        :param connection_observer: The one we are awaiting for.
        :param connection_observer_future: Future of connection-observer returned from submit().
        :return: iterator
        """
        self.logger.debug("go foreground: {!r}".format(connection_observer))

        # assuming that connection_observer.start() / runner.submit(connection_observer)
        # has already scheduled future via asyncio.ensure_future
        assert asyncio.futures.isfuture(connection_observer_future)

        return connection_observer_future.__iter__()
        # Note: even if code is so simple we can't move it inside ConnectionObserver.__await__() since different runners
        # may provide different iterator implementing awaitable
        # Here we know, connection_observer_future is asyncio.Future (precisely asyncio.tasks.Task)
        # and we know it has __await__() method. 

def test_can_receive_binary_data_from_connection(tcp_connection_class,
                                                       integration_tcp_server_and_pipe):
    from moler.threaded_moler_connection import ThreadedMolerConnection
    (tcp_server, tcp_server_pipe) = integration_tcp_server_and_pipe
    received_data = bytearray()
    receiver_called = asyncio.Event()

    def receiver(data, time_recv):
        received_data.extend(data)
        receiver_called.set()

    moler_conn = ThreadedMolerConnection()  # no decoder, just pass bytes 1:1
    moler_conn.subscribe(receiver)       # build forwarding path
    connection = tcp_connection_class(moler_connection=moler_conn, port=tcp_server.port, host=tcp_server.host)
    async with connection:  # TODO: async with connection.open():
        time.sleep(0.1)  # otherwise we have race between server's pipe and from-client-connection
        tcp_server_pipe.send(("send async msg", {'msg': b'data to read'}))
        await asyncio.wait_for(receiver_called.wait(), timeout=0.5)

    assert b'data to read' == received_data


# TODO: tests for error cases raising Exceptions

# --------------------------- resources --------------------------- 

def wait_for(self, event, predicate, result=None):
        """Waits for a DISPATCH'd event that meets the predicate.

        Parameters
        -----------
        event: :class:`str`
            The event name in all upper case to wait for.
        predicate
            A function that takes a data parameter to check for event
            properties. The data parameter is the 'd' key in the JSON message.
        result
            A function that takes the same data parameter and executes to send
            the result to the future. If ``None``, returns the data.

        Returns
        --------
        asyncio.Future
            A future to wait for.
        """

        future = self.loop.create_future()
        entry = EventListener(event=event, predicate=predicate, result=result, future=future)
        self._dispatch_listeners.append(entry)
        return future 

def setUp(self):
        setup_mock_web_api_server(self)

        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        task = asyncio.ensure_future(self.mock_server(), loop=self.loop)
        self.loop.run_until_complete(asyncio.wait_for(task, 0.1))

        self.client = slack.RTMClient(
            token="xoxb-valid",
            base_url="http://localhost:8765",
            auto_reconnect=False,
            run_async=False,
        )
        self.client._web_client = slack.WebClient(
            token="xoxb-valid",
            base_url="http://localhost:8888",
            run_async=False,
        ) 

def run_in_loop(f):
    @functools.wraps(f)
    def wrapper(testcase, *args, **kwargs):
        coro = asyncio.coroutine(f)
        future = asyncio.wait_for(coro(testcase, *args, **kwargs),
                                  timeout=testcase.timeout)
        return testcase.loop.run_until_complete(future)
    return wrapper 

def _execute(
        self, document: DocumentNode, *args, **kwargs
    ) -> ExecutionResult:

        # Fetch schema from transport if needed and validate document if possible
        await self.fetch_and_validate(document)

        # Execute the query with the transport with a timeout
        return await asyncio.wait_for(
            self.transport.execute(document, *args, **kwargs),
            self.client.execute_timeout,
        ) 

def __send_and_receive(self, operation, data):
        async def inner(operation, data):
            if data.request_id is None:
                data.request_id = Util.get_request_id()

            validation = Util.is_invalid(data)
            if validation:
                raise ValidationException(validation)

            # If we're not connected, immediately try to reconnect
            if not self.connected:
                await self.__connect()

            self.__requests[data.request_id] = asyncio.Queue(1)

            # Write request to socket
            frame = MessageFrame(operation=operation, payload=cbor2.dumps(data.as_dict()))
            self.__writer.write(Util.encode_frame(frame))
            await self.__writer.drain()

            # Wait for reader to come back with the response
            result = await self.__requests[data.request_id].get()
            # Drop async queue from request map
            del self.__requests[data.request_id]
            if isinstance(result, MessageFrame) and result.operation == Operation.Unknown:
                raise ClientException("Received response with unknown operation from server")
            return result

        # Perform the actual work as async so that we can put a timeout on the whole operation
        try:
            return await asyncio.wait_for(inner(operation, data), timeout=self.request_timeout, loop=self.__loop)
        except asyncio.TimeoutError:
            # Drop async queue from request map
            del self.__requests[data.request_id]
            raise 

def test_election_many_wait_for_leadership(zk, path):
    NUM = 10
    elec = zk.recipes.LeaderElection(path)

    await elec.volunteer()

    for _ in range(NUM):
        # NO WAIT
        await asyncio.wait_for(elec.wait_for_leadership(), timeout=0.5)
        assert elec.has_leadership

    await elec.resign()
    await zk.delete(path) 

def profile(self, ctx, user: discord.Member = None):
        """Show your leveler progress. Default to yourself."""
        if user is None:
            user = ctx.author
        data = await self.profile_data(user)

        task = functools.partial(self.make_full_profile, **data)
        task = self.bot.loop.run_in_executor(None, task)
        try:
            img = await asyncio.wait_for(task, timeout=60)
        except asyncio.TimeoutError:
            return

        img.seek(0)
        await ctx.send(file=discord.File(img)) 

def start(actual_coroutine):
    """
    Start the testing coroutine and wait 1 second for it to complete.
    :raises asyncio.CancelledError when the coroutine fails to finish its work
        in 1 second.
    :returns: the return value of the actual_coroutine.
    :rtype: Any

    """
    try:
        return await asyncio.wait_for(actual_coroutine, 2)
    except asyncio.CancelledError:
        pass 

def test_inconsistent_zxid():
    async def coro():
        zk = get_client()
        await zk.start()
        # simulate failed connection
        await zk.session.close()
        zk.session.last_zxid = 1231231241312312
        await zk.session.start()
        await zk.session.close()
    try:
        await asyncio.wait_for(coro(), timeout=10)
    except asyncio.TimeoutError as exc:
        pytest.fail("Failed with timeout on session reconnection attemt") 

def test_closed_close():
    zk = get_client()
    await asyncio.wait_for(zk.session.close(), 2) 

def _load_roles(self):
        log.debug(f"Loading roles on {self.guild.id}")
        existing_roles = list(reversed(list(filter(
            lambda r: not r.managed and not r.is_default()
                      and self.guild.me.top_role.position > r.position,
            self.guild.roles
        ))))
        for role in reversed(self.data["roles"]):
            try:
                if role["default"]:
                    await self.guild.default_role.edit(
                        permissions=discord.Permissions(role["permissions"])
                    )
                    new_role = self.guild.default_role
                else:
                    kwargs = {
                        "name": role["name"],
                        "hoist": role["hoist"],
                        "mentionable": role["mentionable"],
                        "color": discord.Color(role["color"]),
                        "permissions": discord.Permissions.none(),
                        "reason": self.reason
                    }

                    if len(existing_roles) == 0:
                        try:
                            new_role = await asyncio.wait_for(self.guild.create_role(**kwargs), 10)
                        except asyncio.TimeoutError:
                            # Probably hit the 24h rate limit. Just skip roles
                            break
                    else:
                        new_role = existing_roles.pop(0)
                        await new_role.edit(**kwargs)

                self.id_translator[role["id"]] = new_role.id
            except Exception:
                pass 

def terminate(self):
        """Terminate a running script."""
        try:
            self.proc.terminate()
        except ProcessLookupError:
            # Process has already been terminated. Log exception and continue.
            logger.exception("While terminating process with PID %s", self.proc.pid)

        await asyncio.wait_for(self.proc.wait(), self.kill_delay)
        if self.proc.returncode is None:
            self.proc.kill()
        await self.proc.wait()

        await super().terminate()