Python trio.fail_after() Examples

def test_peer_packer_sends_who_are_you(peer_packer,
                                             incoming_packet_channels,
                                             outgoing_packet_channels,
                                             nursery):
    incoming_packet = IncomingPacket(
        AuthTagPacketFactory(),
        EndpointFactory(),
    )

    incoming_packet_channels[0].send_nowait(incoming_packet)
    with trio.fail_after(0.5):
        outgoing_packet = await outgoing_packet_channels[1].receive()

    assert peer_packer.is_during_handshake
    assert outgoing_packet.receiver_endpoint == incoming_packet.sender_endpoint
    assert isinstance(outgoing_packet.packet, WhoAreYouPacket)
    assert outgoing_packet.packet.token == incoming_packet.packet.auth_tag 

def heartbeat(ws, timeout, interval):
    '''
    Send periodic pings on WebSocket ``ws``.

    Wait up to ``timeout`` seconds to send a ping and receive a pong. Raises
    ``TooSlowError`` if the timeout is exceeded. If a pong is received, then
    wait ``interval`` seconds before sending the next ping.

    This function runs until cancelled.

    :param ws: A WebSocket to send heartbeat pings on.
    :param float timeout: Timeout in seconds.
    :param float interval: Interval between receiving pong and sending next
        ping, in seconds.
    :raises: ``ConnectionClosed`` if ``ws`` is closed.
    :raises: ``TooSlowError`` if the timeout expires.
    :returns: This function runs until cancelled.
    '''
    while True:
        with trio.fail_after(timeout):
            await ws.ping()
        await trio.sleep(interval) 

def test_datagram_receiver(socket_pair):
    sending_socket, receiving_socket = socket_pair
    receiver_address = receiving_socket.getsockname()
    sender_address = sending_socket.getsockname()

    send_channel, receive_channel = trio.open_memory_channel(1)
    async with background_trio_service(DatagramReceiver(receiving_socket, send_channel)):
        data = b"some packet"

        await sending_socket.sendto(data, receiver_address)
        with trio.fail_after(0.5):
            received_datagram = await receive_channel.receive()

        assert received_datagram.datagram == data
        assert received_datagram.sender_endpoint.ip_address == inet_aton(sender_address[0])
        assert received_datagram.sender_endpoint.port == sender_address[1] 

def test_datagram_sender(socket_pair):
    sending_socket, receiving_socket = socket_pair
    receiver_endpoint = receiving_socket.getsockname()
    sender_endpoint = sending_socket.getsockname()

    send_channel, receive_channel = trio.open_memory_channel(1)
    async with background_trio_service(DatagramSender(receive_channel, sending_socket)):
        outgoing_datagram = OutgoingDatagram(
            b"some packet",
            Endpoint(inet_aton(receiver_endpoint[0]), receiver_endpoint[1]),
        )
        await send_channel.send(outgoing_datagram)

        with trio.fail_after(0.5):
            data, sender = await receiving_socket.recvfrom(1024)
        assert data == outgoing_datagram.datagram
        assert sender == sender_endpoint 

def test_packet_decoder():
    datagram_send_channel, datagram_receive_channel = trio.open_memory_channel(1)
    packet_send_channel, packet_receive_channel = trio.open_memory_channel(1)

    service = PacketDecoder(datagram_receive_channel, packet_send_channel)
    async with background_trio_service(service):
        packet = AuthTagPacketFactory()
        sender_endpoint = EndpointFactory()
        await datagram_send_channel.send(IncomingDatagram(
            datagram=packet.to_wire_bytes(),
            sender_endpoint=sender_endpoint,
        ))

        with trio.fail_after(0.5):
            incoming_packet = await packet_receive_channel.receive()

        assert incoming_packet.packet == packet
        assert incoming_packet.sender_endpoint.ip_address == sender_endpoint.ip_address
        assert incoming_packet.sender_endpoint.port == sender_endpoint.port 

def test_process_while_offline(
    autojump_clock, running_backend, alice_core, bob_user_fs, alice, bob
):
    assert alice_core.backend_status == BackendConnStatus.READY

    with running_backend.offline():
        with alice_core.event_bus.listen() as spy:
            # Force wakeup of the sync monitor
            alice_core.event_bus.send(CoreEvent.FS_ENTRY_UPDATED, id=alice.user_manifest_id)
            assert not alice_core.are_monitors_idle()

            with trio.fail_after(60):  # autojump, so not *really* 60s
                await spy.wait(
                    CoreEvent.BACKEND_CONNECTION_CHANGED,
                    {"status": BackendConnStatus.LOST, "status_exc": spy.ANY},
                )
                await alice_core.wait_idle_monitors()
            assert alice_core.backend_status == BackendConnStatus.LOST 

def test_packet_encoder():
    packet_send_channel, packet_receive_channel = trio.open_memory_channel(1)
    datagram_send_channel, datagram_receive_channel = trio.open_memory_channel(1)

    service = PacketEncoder(packet_receive_channel, datagram_send_channel)
    async with background_trio_service(service):
        receiver_endpoint = EndpointFactory()
        outgoing_packet = OutgoingPacket(
            packet=AuthTagPacketFactory(),
            receiver_endpoint=receiver_endpoint,
        )
        await packet_send_channel.send(outgoing_packet)

        with trio.fail_after(0.5):
            outgoing_datagram = await datagram_receive_channel.receive()

        assert outgoing_datagram.datagram == outgoing_packet.packet.to_wire_bytes()
        assert outgoing_datagram.receiver_endpoint.ip_address == receiver_endpoint.ip_address
        assert outgoing_datagram.receiver_endpoint.port == receiver_endpoint.port 

def test_user_invite_claim_cancel_invitation(monitor, running_backend, backend, alice):
    new_device_id = DeviceID("zack@pc1")
    token = generate_invitation_token()

    invite_and_claim_cancel_scope = None

    async def _from_alice():
        nonlocal invite_and_claim_cancel_scope
        with trio.CancelScope() as invite_and_claim_cancel_scope:
            await invite_and_create_user(alice, new_device_id.user_id, is_admin=False, token=token)

    async def _cancel_invite_and_claim():
        invite_and_claim_cancel_scope.cancel()

    await _invite_and_claim(running_backend, _from_alice, _cancel_invite_and_claim)

    # Now make sure the invitation cannot be used
    with trio.fail_after(1):
        with pytest.raises(InviteClaimError) as exc:
            await claim_user(alice.organization_addr, new_device_id, token=token)
        assert (
            str(exc.value)
            == "Cannot retrieve invitation creator: User `zack` doesn't exist in backend"
        ) 

def test_handle_new_upnp_mapping(manually_driven_discovery, endpoint_server):
    manually_driven_discovery._event_bus = endpoint_server
    manually_driven_discovery.manager.run_daemon_task(
        manually_driven_discovery.handle_new_upnp_mapping)
    assert manually_driven_discovery.this_node.address.ip == '127.0.0.1'
    assert manually_driven_discovery.this_node.enr.sequence_number == 1

    await trio.hazmat.checkpoint()
    external_ip = '43.248.27.0'
    await endpoint_server.wait_until_any_endpoint_subscribed_to(UPnPMapping)
    await endpoint_server.broadcast(UPnPMapping(external_ip))

    with trio.fail_after(0.5):
        while True:
            await trio.sleep(0.01)
            if manually_driven_discovery.this_node.address.ip == external_ip:
                break
    assert manually_driven_discovery.this_node.enr.sequence_number == 2 

def test_concurrent_user_invite(
    backend, apiv1_alice_backend_sock, apiv1_adam_backend_sock, alice, adam, mallory
):
    with backend.event_bus.listen() as spy, trio.fail_after(1):
        async with user_invite(apiv1_alice_backend_sock, user_id=mallory.user_id) as prep1:

            await spy.wait(MetaEvent.EVENT_CONNECTED, {"event_type": BackendEvent.USER_CLAIMED})
            async with user_invite(apiv1_adam_backend_sock, user_id=mallory.user_id) as prep2:

                spy.clear()
                await spy.wait(MetaEvent.EVENT_CONNECTED, {"event_type": BackendEvent.USER_CLAIMED})

                backend.event_bus.send(
                    BackendEvent.USER_CLAIMED,
                    organization_id=mallory.organization_id,
                    user_id=mallory.user_id,
                    encrypted_claim=b"<good>",
                )

    assert prep1[0] == {"status": "ok", "encrypted_claim": b"<good>"}
    assert prep2[0] == {"status": "ok", "encrypted_claim": b"<good>"} 

def test_server_handler_exit(nursery, autojump_clock):
    async def handler(request):
        server_ws = await request.accept()
        await trio.sleep(1)

    server = await nursery.start(
        partial(serve_websocket, handler, HOST, 0, ssl_context=None))

    # connection should close when server handler exists
    with trio.fail_after(2):
        async with open_websocket(
                HOST, server.port, '/', use_ssl=False) as connection:
            with pytest.raises(ConnectionClosed) as exc_info:
                await connection.get_message()
            exc = exc_info.value
            assert exc.reason.name == 'NORMAL_CLOSURE' 

def _start_serving(
        self, nursery: trio_typing.Nursery, ipc_path: pathlib.Path
    ) -> None:
        if not self.is_running:
            raise LifecycleError("Cannot start server if endpoint is not running")
        elif self.is_stopped:
            raise LifecycleError("Endpoint has already been run and stopped")
        elif self.is_serving:
            raise LifecycleError("Endpoint is already serving")
        elif self.is_server_stopped:
            raise LifecycleError("Endpoint server already ran and was stopped")

        self.ipc_path = ipc_path
        nursery.start_soon(self._run_server)

        # Wait until the ipc socket has appeared and is accepting connections.
        with trio.fail_after(constants.IPC_WAIT_SECONDS):
            await _wait_for_path(trio.Path(ipc_path))

        await self._socket_bound.wait() 

def test_peer_packer_initiates_handshake(peer_packer,
                                               outgoing_message_channels,
                                               outgoing_packet_channels,
                                               nursery):
    outgoing_message = OutgoingMessage(
        PingMessageFactory(),
        EndpointFactory(),
        peer_packer.remote_node_id,
    )

    outgoing_message_channels[0].send_nowait(outgoing_message)
    with trio.fail_after(0.5):
        outgoing_packet = await outgoing_packet_channels[1].receive()

    assert peer_packer.is_during_handshake
    assert outgoing_packet.receiver_endpoint == outgoing_message.receiver_endpoint
    assert isinstance(outgoing_packet.packet, AuthTagPacket) 

def test_device_invite_same_name_different_organizations(
    backend, apiv1_alice_backend_sock, apiv1_otheralice_backend_sock, alice, otheralice, alice_nd_id
):
    with backend.event_bus.listen() as spy, trio.fail_after(1):
        async with device_invite(
            apiv1_alice_backend_sock, invited_device_name=alice_nd_id.device_name
        ) as prep:

            await spy.wait(MetaEvent.EVENT_CONNECTED, {"event_type": BackendEvent.DEVICE_CLAIMED})

            backend.event_bus.send(
                BackendEvent.DEVICE_CLAIMED,
                organization_id=otheralice.organization_id,
                device_id=alice_nd_id,
                encrypted_claim=b"<from OtherOrg>",
            )
            await trio.sleep(0)
            backend.event_bus.send(
                BackendEvent.DEVICE_CLAIMED,
                organization_id=alice.organization_id,
                device_id=alice_nd_id,
                encrypted_claim=b"<from CoolOrg>",
            )

    assert prep[0] == {"status": "ok", "encrypted_claim": b"<from CoolOrg>"} 

def test_packer_sends_packets(nursery,
                                    packer,
                                    remote_enr,
                                    remote_endpoint,
                                    outgoing_message_channels,
                                    outgoing_packet_channels):
    assert not packer.is_peer_packer_registered(remote_enr.node_id)

    # send message
    outgoing_message = OutgoingMessage(
        message=PingMessageFactory(),
        receiver_endpoint=remote_endpoint,
        receiver_node_id=remote_enr.node_id,
    )
    outgoing_message_channels[0].send_nowait(outgoing_message)

    with trio.fail_after(0.5):
        outgoing_packet = await outgoing_packet_channels[1].receive()

    assert packer.is_peer_packer_registered(remote_enr.node_id)

    assert isinstance(outgoing_packet.packet, AuthTagPacket)
    assert outgoing_packet.receiver_endpoint == remote_endpoint 

def test_request_handling(message_dispatcher,
                                incoming_message_channels,
                                remote_enr,
                                remote_endpoint):
    ping_send_channel, ping_receive_channel = trio.open_memory_channel(0)

    async with message_dispatcher.add_request_handler(PingMessage) as request_subscription:

        incoming_message = IncomingMessage(
            message=PingMessageFactory(),
            sender_endpoint=remote_endpoint,
            sender_node_id=remote_enr.node_id,
        )
        await incoming_message_channels[0].send(incoming_message)

        with trio.fail_after(1):
            handled_incoming_message = await request_subscription.receive()
        assert handled_incoming_message == incoming_message 

def _wait_for_winfsp_ready(mountpoint_path, timeout=1.0):
    trio_mountpoint_path = trio.Path(mountpoint_path)

    # Polling for `timeout` seconds until winfsp is ready
    with trio.fail_after(timeout):
        while True:
            try:
                if await trio_mountpoint_path.exists():
                    return
                await trio.sleep(0.01)
            # Looks like a revoked workspace has been mounted
            except PermissionError:
                return
            # Might be another OSError like errno 113 (No route to host)
            except OSError:
                return 

def _reql_timeout(seconds):
    """
    Run a block with a timeout, raising `ReqlTimeoutError` if the block
    execution exceeds the timeout.

    :param float seconds: A timeout in seconds. If None, then no timeout is
        enforced.
    :raises ReqlTimeoutError: If execution time exceeds the timeout.
    """
    if seconds is None:
        yield
    else:
        try:
            with trio.fail_after(seconds):
                yield
        except trio.TooSlow:
            raise ReqlTimeoutError() 

def test_response_handling(message_dispatcher, remote_enr, incoming_message_channels):
    request_id = message_dispatcher.get_free_request_id(remote_enr.node_id)
    async with message_dispatcher.add_response_handler(
        remote_enr.node_id,
        request_id,
    ) as response_subscription:

        incoming_message = IncomingMessage(
            message=PingMessageFactory(
                request_id=request_id,
            ),
            sender_endpoint=remote_endpoint,
            sender_node_id=remote_enr.node_id,
        )
        await incoming_message_channels[0].send(incoming_message)

        with trio.fail_after(1):
            handled_response = await response_subscription.receive()
        assert handled_response == incoming_message 

def test_unmount_with_fusermount(base_mountpoint, alice, alice_user_fs, event_bus):
    wid = await alice_user_fs.workspace_create("w")
    workspace = alice_user_fs.get_workspace(wid)
    await workspace.touch("/bar.txt")

    async with mountpoint_manager_factory(
        alice_user_fs, event_bus, base_mountpoint
    ) as mountpoint_manager:

        with event_bus.listen() as spy:
            mountpoint_path = await mountpoint_manager.mount_workspace(wid)
            command = f"fusermount -u {mountpoint_path}".split()
            expected = {"mountpoint": mountpoint_path, "workspace_id": wid, "timestamp": None}

            completed_process = await trio.run_process(command)
            with trio.fail_after(1):
                # fusermount might fail for some reasons
                while completed_process.returncode:
                    completed_process = await trio.run_process(command)
                await spy.wait(CoreEvent.MOUNTPOINT_STOPPED, expected)

        assert not await trio.Path(mountpoint_path / "bar.txt").exists()

    # Mountpoint path should be removed on umounting
    assert not await trio.Path(mountpoint_path).exists() 

def test_init_online_backend_late_reply(
    server_factory, core_config, alice, event_bus, backend
):
    can_serve_client = trio.Event()

    async def _handle_client(stream):
        await can_serve_client.wait()
        return await backend.handle_client(stream)

    async with server_factory(_handle_client, alice.organization_addr):
        with trio.fail_after(1):
            async with logged_core_factory(
                config=core_config, device=alice, event_bus=event_bus
            ) as core:
                # We don't want for backend to reply finish core init
                with core.event_bus.listen() as spy:
                    can_serve_client.set()
                    # Now backend reply, monitor should send events accordingly
                    await spy.wait(
                        CoreEvent.BACKEND_CONNECTION_CHANGED,
                        kwargs={"status": BackendConnStatus.READY, "status_exc": None},
                    ) 

def connect_to_endpoints(self, *endpoints: ConnectionConfig) -> None:
        """
        Connect to the given endpoints and await until all connections are established.
        """
        if not self.is_running:
            raise ConnectionAttemptRejected(
                "Cannot establish connections if endpoint isn't running"
            )

        for config in endpoints:
            # Ensure we are not already connected to the named endpoint
            if self.is_connected_to(config.name):
                raise ConnectionAttemptRejected(
                    f"Already connected to endpoint with name: {config.name}"
                )

            # Feed the `ConnectionConfig` through a channel where
            # `_process_connections` will pick it up and actually establish the
            # connection.
            await self._connection_send_channel.send(config)

        for config in endpoints:
            with trio.fail_after(constants.ENDPOINT_CONNECT_TIMEOUT):
                await self.wait_until_connected_to(config.name)

    #
    # Primary endpoint API
    # 

def check_handshake_timeout(self, handshake_successful_event: trio.Event) -> None:
        try:
            with trio.fail_after(HANDSHAKE_TIMEOUT):
                # Only the timeout for successful handshakes has to be checked as a failure during
                # handshake will make the service as a whole fail.
                await handshake_successful_event.wait()
        except trio.TooSlowError as too_slow_error:
            self.logger.warning("Handshake with %s has timed out", encode_hex(self.remote_node_id))
            raise HandshakeFailure("Handshake has timed out") from too_slow_error

    #
    # Handshake states
    # 

def test_send_endpoint_vote(ping_sender_service,
                                  routing_table,
                                  incoming_message_channels,
                                  outgoing_message_channels,
                                  endpoint_vote_channels,
                                  local_enr,
                                  remote_enr,
                                  remote_endpoint):
    # wait for ping
    with trio.fail_after(ROUTING_TABLE_PING_INTERVAL):
        outgoing_message = await outgoing_message_channels[1].receive()
    ping = outgoing_message.message

    # respond with pong
    fake_local_endpoint = EndpointFactory()
    pong = PongMessage(
        request_id=ping.request_id,
        enr_seq=0,
        packet_ip=fake_local_endpoint.ip_address,
        packet_port=fake_local_endpoint.port,
    )
    incoming_message = IncomingMessage(
        message=pong,
        sender_endpoint=outgoing_message.receiver_endpoint,
        sender_node_id=outgoing_message.receiver_node_id,
    )
    await incoming_message_channels[0].send(incoming_message)
    await wait_all_tasks_blocked()

    # receive endpoint vote
    endpoint_vote = endpoint_vote_channels[1].receive_nowait()
    assert endpoint_vote.endpoint == fake_local_endpoint
    assert endpoint_vote.node_id == incoming_message.sender_node_id 

def request_enr(self, remote: NodeAPI) -> ENR:
        """Get the most recent ENR for the given node and update our local DB and routing table.

        The updating of the DB and RT happens in the handler called when we receive an ENR
        response, so that the new ENR is stored even if we give up waiting.

        Raises CouldNotRetrieveENR if we can't get a ENR from the remote node.
        """
        # No need to use a timeout because bond() takes care of that internally.
        await self.bond(remote.id)
        token = self.send_enr_request(remote)
        send_chan, recv_chan = trio.open_memory_channel[Tuple[ENR, Hash32]](1)
        try:
            with trio.fail_after(constants.KADEMLIA_REQUEST_TIMEOUT):
                enr, received_token = await self.enr_response_channels.receive_one(
                    remote, send_chan, recv_chan)
        except trio.TooSlowError:
            raise CouldNotRetrieveENR(f"Timed out waiting for ENR from {remote}")
        except AlreadyWaitingDiscoveryResponse:
            raise CouldNotRetrieveENR(f"Already waiting for ENR from {remote}")

        if received_token != token:
            raise CouldNotRetrieveENR(
                f"Got ENR from {remote} with token {received_token!r} but expected {token!r}")

        return enr 

def test_trio_server_endpoint_establishes_reverse_connection_to_client(
    ipc_base_path
):
    unique_name = generate_unique_name()
    config = ConnectionConfig.from_name(
        f"server-{unique_name}", base_path=ipc_base_path
    )

    async with TrioEndpoint.serve(config) as server:
        async with TrioEndpoint(f"client-{unique_name}").run() as client:
            await client.connect_to_endpoints(config)

            assert client.is_connected_to(config.name)
            with trio.fail_after(2):
                await server.wait_until_connected_to(client.name) 

def test_get_peer_candidates(manually_driven_discovery, monkeypatch):
    total_nodes = 10
    nodes = NodeFactory.create_batch(total_nodes)
    discovery = manually_driven_discovery
    for node in nodes:
        discovery.node_db.set_enr(node.enr)
        assert discovery.routing.update(node.id) is None

    discovery._random_lookup_calls = 0

    async def mock_lookup_random():
        discovery._random_lookup_calls += 1

    monkeypatch.setattr(discovery, 'lookup_random', mock_lookup_random)

    def should_skip(skip_list, candidate):
        return candidate in skip_list

    candidates = discovery.get_peer_candidates(functools.partial(should_skip, tuple()), total_nodes)
    assert sorted(candidates) == sorted(nodes)

    candidates = discovery.get_peer_candidates(
        functools.partial(should_skip, tuple()), total_nodes + 10)
    assert sorted(candidates) == sorted(nodes)
    # When we don't have enough candidates, a random lookup should be triggered.
    with trio.fail_after(0.5):
        while discovery._random_lookup_calls != 1:
            await trio.sleep(0.01)

    candidates = discovery.get_peer_candidates(
        functools.partial(should_skip, tuple()), total_nodes - 1)
    assert len(candidates) == total_nodes - 1

    skip_list = (nodes[0], nodes[5], nodes[8])
    candidates = discovery.get_peer_candidates(
        functools.partial(should_skip, skip_list), total_nodes)
    assert sorted(candidates) == sorted(set(nodes).difference(skip_list))
    with trio.fail_after(0.5):
        while discovery._random_lookup_calls != 2:
            await trio.sleep(0.01) 

def _process_connections(
        self,
        channel: trio.abc.ReceiveChannel[ConnectionConfig],
        nursery: trio_typing.Nursery,
    ) -> None:
        """
        Long running process that establishes connections to endpoint servers
        and runs the handler for receiving events sent by the server over that
        connection.
        """
        self.logger.debug("%s: starting new connection channel", self)
        self._connection_loop_running.set()
        async for config in channel:
            # Allow some time for for the IPC socket to appear
            with trio.fail_after(constants.IPC_WAIT_SECONDS):
                await _wait_for_path(trio.Path(config.path))

            await trio.sleep(0.001)
            # Establish the socket connection
            connection = await TrioConnection.connect_to(config.path)

            # Create the remote
            remote = TrioRemoteEndpoint(
                self.name,
                connection,
                self._remote_subscriptions_changed,
                self._inbound_send_channel.send,
            )
            nursery.start_soon(self._run_remote_endpoint, remote) 

def _start(self, nursery: trio_typing.Nursery) -> None:
        if self.is_running:
            raise LifecycleError("Endpoint is already running")
        elif self.is_stopped:
            raise LifecycleError("Endpoint has already been run and stopped")

        with trio.fail_after(2):
            await nursery.start(self._run)

            # this is a trade-off between allowing flexibility in how long it takes
            # for a `Runnable` to start and getting good errors in the event that
            # the `_run` method forgets to set the `_running` event.
            await self.wait_started() 

def _start(self, nursery: trio_typing.Nursery) -> None:
        if self.is_running:
            raise LifecycleError("RemoteEndpoint is already running")
        elif self.is_stopped:
            raise LifecycleError("RemoteEndpoint has already been run and stopped")

        nursery.start_soon(self._process_incoming_messages)
        # this is a trade-off between allowing flexibility in how long it takes
        # for a `Runnable` to start and getting good errors in the event that
        # the `_run` method forgets to set the `_running` event.
        with trio.fail_after(2):
            await self.wait_started()