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

Code before:
            p.send_message(msg_inv([CInv(t=MSG_WTX, h=WTXID)]))
        # One of the peers is asked for the tx
        peer2.wait_until(lambda: sum(p.tx_getdata_count for p in [peer1, peer2]) == 1)
        with p2p_lock:
            peer_disconnect, peer_fallback = (peer1, peer2) if peer1.tx_getdata_count == 1 else (peer2, peer1)
            assert_equal(peer_fallback.tx_getdata_count, 0)
        peer_disconnect.peer_disconnect()
        peer_disconnect.wait_for_disconnect()
        peer_fallback.wait_until(lambda: peer_fallback.tx_getdata_count >= 1, timeout=1)

    def test_notfound_fallback(self):
        self.log.info('Check that notfounds will select another peer for download immediately')
        WTXID = 0xffdd
        peer1 = self.nodes[0].add_p2p_connection(TestP2PConn())
        peer2 = self.nodes[0].add_p2p_connection(TestP2PConn())
        for p in [peer1, peer2]:
            p.send_message(msg_inv([CInv(t=MSG_WTX, h=WTXID)]))
        # One of the peers is asked for the tx
        peer2.wait_until(lambda: sum(p.tx_getdata_count for p in [peer1, peer2]) == 1)
        with p2p_lock:
            peer_notfound, peer_fallback = (peer1, peer2) if peer1.tx_getdata_count == 1 else (peer2, peer1)
            assert_equal(peer_fallback.tx_getdata_count, 0)
        peer_notfound.send_and_ping(msg_notfound(vec=[CInv(MSG_WTX, WTXID)]))  # Send notfound, so that fallback peer is selected
        peer_fallback.wait_until(lambda: peer_fallback.tx_getdata_count >= 1, timeout=1)

    def test_preferred_inv(self, preferred=False):
        if preferred:
            self.log.info('Check invs from preferred peers are downloaded immediately')
            self.restart_node(0, extra_args=['-whitelist=noban@127.0.0.1'])
        else:
            self.log.info('Check invs from non-preferred peers are downloaded after {} s'.format(NONPREF_PEER_TX_DELAY))
        mock_time = int(time.time() + 1)
        self.nodes[0].setmocktime(mock_time)
        peer = self.nodes[0].add_p2p_connection(TestP2PConn())
        peer.send_message(msg_inv([CInv(t=MSG_WTX, h=0xff00ff00)]))
        peer.sync_with_ping()
        if preferred:
            peer.wait_until(lambda: peer.tx_getdata_count >= 1, timeout=1)
        else:
            with p2p_lock:
                assert_equal(peer.tx_getdata_count, 0)
            self.nodes[0].setmocktime(mock_time + NONPREF_PEER_TX_DELAY)
            peer.wait_until(lambda: peer.tx_getdata_count >= 1, timeout=1)

    def test_txid_inv_delay(self, glob_wtxid=False):
        self.log.info('Check that inv from a txid-relay peers are delayed by {} s, with a wtxid peer {}'.format(TXID_RELAY_DELAY, glob_wtxid))
        self.restart_node(0, extra_args=['-whitelist=noban@127.0.0.1'])
        mock_time = int(time.time() + 1)
        self.nodes[0].setmocktime(mock_time)
        peer = self.nodes[0].add_p2p_connection(TestP2PConn(wtxidrelay=False))
        if glob_wtxid:

Review comment (suggestion):
nit: consider updating the preferred_type with a type hint and an enum to avoid using "magic strings"

Code after:
        self.nodes[0].setmocktime(int(time.time()) + GETDATA_TX_INTERVAL + 1)  # Wait for request to _peer_expiry to expire
        peer_fallback.wait_until(lambda: peer_fallback.tx_getdata_count >= 1, timeout=1)
        self.restart_node(0)  # reset mocktime

    def test_disconnect_fallback(self):
        self.log.info('Check that disconnect will select another peer for download')
        WTXID = 0xffbb
        peer1 = self.nodes[0].add_p2p_connection(TestP2PConn())
        peer2 = self.nodes[0].add_p2p_connection(TestP2PConn())
        for p in [peer1, peer2]:
            p.send_message(msg_inv([CInv(t=MSG_WTX, h=WTXID)]))
        # One of the peers is asked for the tx
        peer2.wait_until(lambda: sum(p.tx_getdata_count for p in [peer1, peer2]) == 1)
        with p2p_lock:
            peer_disconnect, peer_fallback = (peer1, peer2) if peer1.tx_getdata_count == 1 else (peer2, peer1)
            assert_equal(peer_fallback.tx_getdata_count, 0)
        peer_disconnect.peer_disconnect()
        peer_disconnect.wait_for_disconnect()
        peer_fallback.wait_until(lambda: peer_fallback.tx_getdata_count >= 1, timeout=1)

    def test_notfound_fallback(self):
        self.log.info('Check that notfounds will select another peer for download immediately')
        WTXID = 0xffdd
        peer1 = self.nodes[0].add_p2p_connection(TestP2PConn())
        peer2 = self.nodes[0].add_p2p_connection(TestP2PConn())
        for p in [peer1, peer2]:
            p.send_message(msg_inv([CInv(t=MSG_WTX, h=WTXID)]))
        # One of the peers is asked for the tx
        peer2.wait_until(lambda: sum(p.tx_getdata_count for p in [peer1, peer2]) == 1)
        with p2p_lock:
            peer_notfound, peer_fallback = (peer1, peer2) if peer1.tx_getdata_count == 1 else (peer2, peer1)
            assert_equal(peer_fallback.tx_getdata_count, 0)
        peer_notfound.send_and_ping(msg_notfound(vec=[CInv(MSG_WTX, WTXID)]))  # Send notfound, so that fallback peer is selected
        peer_fallback.wait_until(lambda: peer_fallback.tx_getdata_count >= 1, timeout=1)

    def test_preferred_inv(self, connection_type: ConnectionType):
        if connection_type == ConnectionType.WHITELIST:
            self.log.info('Check invs from preferred (whitelisted) peers are downloaded immediately')
            self.restart_node(0, extra_args=['-whitelist=noban@127.0.0.1'])
        elif connection_type == ConnectionType.OUTBOUND:
            self.log.info('Check invs from preferred (outbound) peers are downloaded immediately')
            self.restart_node(0)
        elif connection_type == ConnectionType.INBOUND:
            self.log.info('Check invs from non-preferred peers are downloaded after {} s'.format(NONPREF_PEER_TX_DELAY))
            self.restart_node(0)
        else:
            raise Exception("invalid connection_type")

        mock_time = int(time.time() + 1)
        self.nodes[0].setmocktime(mock_time)