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

Code before:
            # The tricky part of this test is the timestamp of the self-announcement.
            # The following helped me to debug timestamp mismatches, so I'm leaving it in for the next one.
            elif addr.ip == self.expected.ip and addr.port == self.expected.port and addr.nServices == self.expected.nServices:
                print(f"Self-announcement timestamp mismatch: got={addr.time} expected={self.expected.time}")

    def on_addrv2(self, message):
        assert (self.addrv2_test)
        self.handle_addr_message(message)

    def on_addr(self, message):
        assert (not self.addrv2_test)
        self.handle_addr_message(message)


class AddrSelfAnnouncementTest(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [[f"-externalip={IP_TO_ANNOUNCE}"]]

    def run_test(self):
        # populate addrman to have some addresses for a GETADDR response
        for i in range(50):
            a = f"{1 + i}.{i}.1.1"
            self.nodes[0].addpeeraddress(a, 8333)

        self.self_announcement_test(outbound=False, addrv2=False, assert_on_connection_open=self.inbound_connection_open_assertions)
        self.self_announcement_test(outbound=False, addrv2=True, assert_on_connection_open=self.inbound_connection_open_assertions)
        self.self_announcement_test(outbound=True, addrv2=False, assert_on_connection_open=self.outbound_connection_open_assertions)
        self.self_announcement_test(outbound=True, addrv2=True, assert_on_connection_open=self.outbound_connection_open_assertions)

    def inbound_connection_open_assertions(self, addr_receiver):
        # We expect one self-announcement and multiple other addresses in
        # response to a GETADDR in a single addr / addrv2 message.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_greater_than(addr_receiver.addresses_received, 1)

    def outbound_connection_open_assertions(self, addr_receiver):
        # We expect only the self-announcement.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_equal(addr_receiver.addresses_received, 1)

    def self_announcement_test(self, outbound, addrv2, assert_on_connection_open):
        connection_type = "outbound" if outbound else "inbound"
        addr_version = "addrv2" if addrv2 else "addrv1"
        self.log.info(f"Test that the node does an address self-announcement to {connection_type} connections ({addr_version})")

        # We only self-announce after initial block download is done
        assert (not self.nodes[0].getblockchaininfo()["initialblockdownload"])


Review comment (suggestion):
`assert_on_connection_open` seems like a redundant arg, could call one or the other based on the existing `outbound` arg.

Code after:

    def on_addrv2(self, message):
        assert (self.addrv2_test)
        self.handle_addr_message(message)

    def on_addr(self, message):
        assert (not self.addrv2_test)
        self.handle_addr_message(message)


class AddrSelfAnnouncementTest(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.extra_args = [[f"-externalip={IP_TO_ANNOUNCE}"]]

    def run_test(self):
        # populate addrman to have some addresses for a GETADDR response
        for i in range(50):
            a = f"{1 + i}.{i}.1.1"
            self.nodes[0].addpeeraddress(a, 8333)

        self.self_announcement_test(outbound=False, addrv2=False)
        self.self_announcement_test(outbound=False, addrv2=True)
        self.self_announcement_test(outbound=True, addrv2=False)
        self.self_announcement_test(outbound=True, addrv2=True)

    def inbound_connection_open_assertions(self, addr_receiver):
        # We expect one self-announcement and multiple other addresses in
        # response to a GETADDR in a single addr / addrv2 message.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_greater_than(addr_receiver.addresses_received, 1)

    def outbound_connection_open_assertions(self, addr_receiver):
        # We expect only the self-announcement.
        assert_equal(addr_receiver.self_announcements_received, 1)
        assert_equal(addr_receiver.addr_messages_received, 1)
        assert_equal(addr_receiver.addresses_received, 1)

    def self_announcement_test(self, *, outbound, addrv2):
        connection_type = "outbound" if outbound else "inbound"
        addr_version = "addrv2" if addrv2 else "addrv1"
        self.log.info(f"Test that the node does an address self-announcement to {connection_type} connections ({addr_version})")

        # We only self-announce after initial block download is done
        assert (not self.nodes[0].getblockchaininfo()["initialblockdownload"])

        netinfo = self.nodes[0].getnetworkinfo()
        port = netinfo["localaddresses"][0]["port"]
        self.nodes[0].setmocktime(int(time.time()))