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

Code before:

IP_TO_ANNOUNCE = "42.42.42.42"
ONE_DAY = 60 * 60 * 24


class SelfAnnouncementReceiver(P2PInterface):
    self_announcements_received = 0
    addresses_received = 0
    addr_messages_received = 0

    expected = None
    addrv2_test = False

    def __init__(self, *, expected, addrv2):
        super().__init__(support_addrv2=addrv2)
        self.expected = expected
        self.addrv2_test = addrv2

    def handle_addr_message(self, message):
        self.addr_messages_received += 1
        for addr in message.addrs:
            self.addresses_received += 1
            if addr == self.expected:
                self.self_announcements_received += 1
            # 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)

Review comment (refactor):
I would suggest to rework or remove this before merge, first person language and `print()` don't seem appropriate at that stage anymore, the mismatch print could be an assert if it's not expected to ever happen.

Code after:

IP_TO_ANNOUNCE = "42.42.42.42"
ONE_DAY = 60 * 60 * 24


class SelfAnnouncementReceiver(P2PInterface):
    self_announcements_received = 0
    addresses_received = 0
    addr_messages_received = 0

    expected = None
    addrv2_test = False

    def __init__(self, *, expected, addrv2):
        super().__init__(support_addrv2=addrv2)
        self.expected = expected
        self.addrv2_test = addrv2

    def handle_addr_message(self, message):
        self.addr_messages_received += 1
        for addr in message.addrs:
            self.addresses_received += 1
            if addr == self.expected:
                self.self_announcements_received += 1

    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):