test: makes wait_for_getdata pops data on checks, plus check the mess…

…age types if provided

Currently, `wait_for_getdata` checks data in `last_message`, but it leaves is as is. This results
in multiple tests having to manually pop data from the collection is multiple getdata messages need
to be checked in order to make sure that we are not asserting the oldest message. This was partially
solved by checking the hash of what was being waited for, but the issue remains if the same tx/block
is sent multiple times during the test.

In some of this scenarios, we were also manually checking the getdata message type manually which is not
possible anymore if we pop data on checking, so `wait_for_getdata` is extended to support checking the type
when we care about it

test/functional/p2p_compactblocks.py

@@ -391,8 +391,7 @@ def test_compactblock_requests(self, test_node):
                 test_node.send_header_for_blocks([block])
             else:
                 test_node.send_header_for_blocks([block])
-            test_node.wait_for_getdata([block.sha256], timeout=30)
-            assert_equal(test_node.last_message["getdata"].inv[0].type, 4)
+            test_node.wait_for_getdata([block.sha256], datatype=MSG_CMPCT_BLOCK, timeout=30)
 
             # Send back a compactblock message that omits the coinbase
             comp_block = HeaderAndShortIDs()
@@ -558,9 +557,7 @@ def test_incorrect_blocktxn_response(self, test_node):
         assert_equal(int(node.getbestblockhash(), 16), block.hashPrevBlock)
 
         # We should receive a getdata request
-        test_node.wait_for_getdata([block.sha256], timeout=10)
-        assert test_node.last_message["getdata"].inv[0].type == MSG_BLOCK or \
-               test_node.last_message["getdata"].inv[0].type == MSG_BLOCK | MSG_WITNESS_FLAG
+        test_node.wait_for_getdata([block.sha256], datatype=(MSG_BLOCK | MSG_WITNESS_FLAG), timeout=10)
 
         # Deliver the block
         test_node.send_and_ping(msg_block(block))

test/functional/p2p_segwit.py

@@ -167,12 +167,10 @@ def on_getdata(self, message):
     def on_wtxidrelay(self, message):
         self.last_wtxidrelay.append(message)
 
-    def announce_tx_and_wait_for_getdata(self, tx, success=True, use_wtxid=False):
+    def announce_tx_and_wait_for_getdata(self, tx, success=True, use_wtxid=False, datatype=None):
         if success:
             # sanity check
             assert (self.wtxidrelay and use_wtxid) or (not self.wtxidrelay and not use_wtxid)
-        with p2p_lock:
-            self.last_message.pop("getdata", None)
         if use_wtxid:
             wtxid = tx.calc_sha256(True)
             self.send_message(msg_inv(inv=[CInv(MSG_WTX, wtxid)]))
@@ -181,16 +179,14 @@ def announce_tx_and_wait_for_getdata(self, tx, success=True, use_wtxid=False):
 
         if success:
             if use_wtxid:
-                self.wait_for_getdata([wtxid])
+                self.wait_for_getdata([wtxid], datatype=datatype)
             else:
-                self.wait_for_getdata([tx.sha256])
+                self.wait_for_getdata([tx.sha256], datatype=datatype)
         else:
             time.sleep(5)
             assert not self.last_message.get("getdata")
 
     def announce_block_and_wait_for_getdata(self, block, use_header, timeout=60):
-        with p2p_lock:
-            self.last_message.pop("getheaders", None)
         msg = msg_headers()
         msg.headers = [CBlockHeader(block)]
         if use_header:
@@ -199,7 +195,7 @@ def announce_block_and_wait_for_getdata(self, block, use_header, timeout=60):
             self.send_message(msg_inv(inv=[CInv(MSG_BLOCK, block.sha256)]))
             self.wait_for_getheaders(timeout=timeout)
             self.send_message(msg)
-        self.wait_for_getdata([block.sha256], timeout=timeout)
+        self.wait_for_getdata([block.sha256], datatype=(MSG_BLOCK | MSG_WITNESS_FLAG), timeout=timeout)
 
     def request_block(self, blockhash, inv_type, timeout=60):
         with p2p_lock:
@@ -361,8 +357,6 @@ def test_block_relay(self):
         This is true regardless of segwit activation.
         Also test that we don't ask for blocks from unupgraded peers."""
 
-        blocktype = 2 | MSG_WITNESS_FLAG
-
         # test_node has set NODE_WITNESS, so all getdata requests should be for
         # witness blocks.
         # Test announcing a block via inv results in a getdata, and that
@@ -375,14 +369,12 @@ def test_block_relay(self):
         self.test_node.send_message(msg_headers())
 
         self.test_node.announce_block_and_wait_for_getdata(block1, use_header=False)
-        assert self.test_node.last_message["getdata"].inv[0].type == blocktype
         test_witness_block(self.nodes[0], self.test_node, block1, True)
 
         block2 = self.build_next_block()
         block2.solve()
 
         self.test_node.announce_block_and_wait_for_getdata(block2, use_header=True)
-        assert self.test_node.last_message["getdata"].inv[0].type == blocktype
         test_witness_block(self.nodes[0], self.test_node, block2, True)
 
         # Check that we can getdata for witness blocks or regular blocks,
@@ -537,8 +529,7 @@ def test_witness_tx_relay_before_segwit_activation(self):
 
         # Verify that if a peer doesn't set nServices to include NODE_WITNESS,
         # the getdata is just for the non-witness portion.
-        self.old_node.announce_tx_and_wait_for_getdata(tx)
-        assert self.old_node.last_message["getdata"].inv[0].type == MSG_TX
+        self.old_node.announce_tx_and_wait_for_getdata(tx, datatype=MSG_TX)
 
         # Since we haven't delivered the tx yet, inv'ing the same tx from
         # a witness transaction ought not result in a getdata.

test/functional/test_framework/p2p.py

@@ -600,12 +600,18 @@ def test_function():
 
     # Message receiving helper methods
 
-    def check_last_getdata(self, hash_list):
-        with p2p_lock:
-            last_data = self.last_message.get("getdata")
-            if last_data is None:
-                return False
+    def _check_last_getdata(self, hash_list, datatype=None):
+        last_data = self.last_message.pop("getdata", None)
+        if not last_data:
+            return False
+        if datatype is None:
             return [x.hash for x in last_data.inv] == hash_list
+        else:
+            return [(x.hash, x.type) for x in last_data.inv] == [(h, datatype) for h in hash_list]
+
+    def check_last_getdata(self, hash_list, datatype=None):
+        with p2p_lock:
+            self._check_last_getdata(hash_list, datatype)
 
     def wait_for_tx(self, txid, *, timeout=60):
         def test_function():
@@ -639,15 +645,12 @@ def test_function():
 
         self.wait_until(test_function, timeout=timeout)
 
-    def wait_for_getdata(self, hash_list, *, timeout=60):
+    def wait_for_getdata(self, hash_list, *, datatype=None, timeout=60):
         """Waits for a getdata message.
 
         The object hashes in the inventory vector must match the provided hash_list."""
         def test_function():
-            last_data = self.last_message.get("getdata")
-            if not last_data:
-                return False
-            return [x.hash for x in last_data.inv] == hash_list
+            return self._check_last_getdata(hash_list, datatype)
 
         self.wait_until(test_function, timeout=timeout)