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

Code before:
    def on_verack(self, message):
        pass

    def on_version(self, message):
        assert message.nVersion >= MIN_P2P_VERSION_SUPPORTED, "Version {} received. Test framework only supports versions greater than {}".format(message.nVersion, MIN_P2P_VERSION_SUPPORTED)
        # for inbound connections, reply to version with own version message
        # (could be due to v1 reconnect after a failed v2 handshake)
        if not self.p2p_connected_to_node:
            self.send_version()
            self.reconnect = False
        if message.nVersion >= 70016 and self.wtxidrelay:
            self.send_message(msg_wtxidrelay())
        if self.support_addrv2:
            self.send_message(msg_sendaddrv2())
        self.send_message(msg_verack())
        self.nServices = message.nServices
        self.relay = message.relay
        if self.p2p_connected_to_node:
            self.send_message(msg_getaddr())

    # Connection helper methods

    def wait_until(self, test_function_in, *, timeout=60, check_connected=True):
        def test_function():
            if check_connected:
                assert self.is_connected
            return test_function_in()

        wait_until_helper_internal(test_function, timeout=timeout, lock=p2p_lock, timeout_factor=self.timeout_factor)

    def wait_for_connect(self, *, timeout=60):
        test_function = lambda: self.is_connected
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    def wait_for_disconnect(self, *, timeout=60):
        test_function = lambda: not self.is_connected
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    def wait_for_reconnect(self, *, timeout=60):
        def test_function():
            return self.is_connected and self.last_message.get('version') and not self.supports_v2_p2p
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    # Message receiving helper methods

    def wait_for_tx(self, txid, *, timeout=60):
        def test_function():
            if not self.last_message.get('tx'):
                return False
            return self.last_message['tx'].tx.rehash() == txid


Review comment (suggestion):
```suggestion
    def ensure(self, test_function, *, **kwargs):
```

I am not sure about a default value here. I feel like `time.sleep()`, or similar constructs in tests should be easily visible. Otherwise, this may be used as a drop-in replacement for `assert`, because someone may think this is just an assert that accepts a lambda.

Code after:
        # (could be due to v1 reconnect after a failed v2 handshake)
        if not self.p2p_connected_to_node:
            self.send_version()
            self.reconnect = False
        if message.nVersion >= 70016 and self.wtxidrelay:
            self.send_message(msg_wtxidrelay())
        if self.support_addrv2:
            self.send_message(msg_sendaddrv2())
        self.send_message(msg_verack())
        self.nServices = message.nServices
        self.relay = message.relay
        if self.p2p_connected_to_node:
            self.send_message(msg_getaddr())

    # Connection helper methods

    def wait_until(self, test_function_in, *, timeout=60, check_connected=True, check_interval=0.05):
        def test_function():
            if check_connected:
                assert self.is_connected
            return test_function_in()

        wait_until_helper_internal(test_function, timeout=timeout, lock=p2p_lock, timeout_factor=self.timeout_factor, check_interval=check_interval)

    def wait_for_connect(self, *, timeout=60):
        test_function = lambda: self.is_connected
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    def wait_for_disconnect(self, *, timeout=60):
        test_function = lambda: not self.is_connected
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    def wait_for_reconnect(self, *, timeout=60):
        def test_function():
            return self.is_connected and self.last_message.get('version') and not self.supports_v2_p2p
        self.wait_until(test_function, timeout=timeout, check_connected=False)

    # Message receiving helper methods

    def wait_for_tx(self, txid, *, timeout=60):
        def test_function():
            if not self.last_message.get('tx'):
                return False
            return self.last_message['tx'].tx.rehash() == txid

        self.wait_until(test_function, timeout=timeout)

    def wait_for_block(self, blockhash, *, timeout=60):
        def test_function():
            return self.last_message.get("block") and self.last_message["block"].block.rehash() == blockhash