Repository: bitcoin/bitcoin
File: test/functional/test_framework/test_framework.py (Python)

Code before:
            # compatibility with older clients
            from_connection.addnode(ip_port, "onetry")

        if not wait_for_connect:
            return

        # poll until version handshake complete to avoid race conditions
        # with transaction relaying
        # See comments in net_processing:
        # * Must have a version message before anything else
        # * Must have a verack message before anything else
        self.wait_until(lambda: sum(peer['version'] != 0 for peer in from_connection.getpeerinfo()) == from_num_peers)
        self.wait_until(lambda: sum(peer['version'] != 0 for peer in to_connection.getpeerinfo()) == to_num_peers)
        self.wait_until(lambda: sum(peer['bytesrecv_per_msg'].pop('verack', 0) >= 21 for peer in from_connection.getpeerinfo()) == from_num_peers)
        self.wait_until(lambda: sum(peer['bytesrecv_per_msg'].pop('verack', 0) >= 21 for peer in to_connection.getpeerinfo()) == to_num_peers)
        # The message bytes are counted before processing the message, so make
        # sure it was fully processed by waiting for a ping.
        self.wait_until(lambda: sum(peer["bytesrecv_per_msg"].pop("pong", 0) >= 29 for peer in from_connection.getpeerinfo()) == from_num_peers)
        self.wait_until(lambda: sum(peer["bytesrecv_per_msg"].pop("pong", 0) >= 29 for peer in to_connection.getpeerinfo()) == to_num_peers)

    def disconnect_nodes(self, a, b):
        def disconnect_nodes_helper(node_a, node_b):
            def get_peer_ids(from_connection, node_num):
                result = []
                for peer in from_connection.getpeerinfo():
                    if "testnode{}".format(node_num) in peer['subver']:
                        result.append(peer['id'])
                return result

            peer_ids = get_peer_ids(node_a, node_b.index)
            if not peer_ids:
                self.log.warning("disconnect_nodes: {} and {} were not connected".format(
                    node_a.index,
                    node_b.index,
                ))
                return
            for peer_id in peer_ids:
                try:
                    node_a.disconnectnode(nodeid=peer_id)
                except JSONRPCException as e:
                    # If this node is disconnected between calculating the peer id
                    # and issuing the disconnect, don't worry about it.
                    # This avoids a race condition if we're mass-disconnecting peers.
                    if e.error['code'] != -29:  # RPC_CLIENT_NODE_NOT_CONNECTED
                        raise

            # wait to disconnect
            self.wait_until(lambda: not get_peer_ids(node_a, node_b.index), timeout=5)
            self.wait_until(lambda: not get_peer_ids(node_b, node_a.index), timeout=5)

        disconnect_nodes_helper(self.nodes[a], self.nodes[b])

Review comment (suggestion):
`find_conn(from_connection, to_connection_subver, inbound=False)`

Might as well make these calls once and store in variable instead of finding 3 times? Unless I'm missing something that requires these calls to be made every time. 


Code after:

        if not wait_for_connect:
            return

        # Use subversion as peer id. Test nodes have their node number appended to the user agent string
        from_connection_subver = from_connection.getnetworkinfo()['subversion']
        to_connection_subver = to_connection.getnetworkinfo()['subversion']

        def find_conn(node, peer_subversion, inbound):
            return next(filter(lambda peer: peer['subver'] == peer_subversion and peer['inbound'] == inbound, node.getpeerinfo()), None)

        # poll until version handshake complete to avoid race conditions
        # with transaction relaying
        # See comments in net_processing:
        # * Must have a version message before anything else
        # * Must have a verack message before anything else
        self.wait_until(lambda: find_conn(from_connection, to_connection_subver, inbound=False) is not None)
        self.wait_until(lambda: find_conn(to_connection, from_connection_subver, inbound=True) is not None)

        def check_bytesrecv(peer, msg_type, min_bytes_recv):
            assert peer is not None, "Error: peer disconnected"
            return peer['bytesrecv_per_msg'].pop(msg_type, 0) >= min_bytes_recv

        self.wait_until(lambda: check_bytesrecv(find_conn(from_connection, to_connection_subver, inbound=False), 'verack', 21))
        self.wait_until(lambda: check_bytesrecv(find_conn(to_connection, from_connection_subver, inbound=True), 'verack', 21))

        # The message bytes are counted before processing the message, so make
        # sure it was fully processed by waiting for a ping.
        self.wait_until(lambda: check_bytesrecv(find_conn(from_connection, to_connection_subver, inbound=False), 'pong', 29))
        self.wait_until(lambda: check_bytesrecv(find_conn(to_connection, from_connection_subver, inbound=True), 'pong', 29))

    def disconnect_nodes(self, a, b):
        def disconnect_nodes_helper(node_a, node_b):
            def get_peer_ids(from_connection, node_num):
                result = []
                for peer in from_connection.getpeerinfo():
                    if "testnode{}".format(node_num) in peer['subver']:
                        result.append(peer['id'])
                return result

            peer_ids = get_peer_ids(node_a, node_b.index)
            if not peer_ids:
                self.log.warning("disconnect_nodes: {} and {} were not connected".format(
                    node_a.index,
                    node_b.index,
                ))
                return
            for peer_id in peer_ids:
                try:
                    node_a.disconnectnode(nodeid=peer_id)
                except JSONRPCException as e: