[tests] fix wait_for_inv()

Author: jnewbery

test/functional/p2p-segwit.py

@@ -926,9 +926,9 @@ def test_tx_relay_after_segwit_activation(self):
         tx3.wit.vtxinwit[0].scriptWitness.stack = [ witness_program ]
         # Also check that old_node gets a tx announcement, even though this is
         # a witness transaction.
-        self.old_node.wait_for_inv(CInv(1, tx2.sha256)) # wait until tx2 was inv'ed
+        self.old_node.wait_for_inv([CInv(1, tx2.sha256)]) # wait until tx2 was inv'ed
         self.test_node.test_transaction_acceptance(tx3, with_witness=True, accepted=True)
-        self.old_node.wait_for_inv(CInv(1, tx3.sha256))
+        self.old_node.wait_for_inv([CInv(1, tx3.sha256)])
 
         # Test that getrawtransaction returns correct witness information
         # hash, size, vsize

test/functional/test_framework/mininode.py

@@ -1604,7 +1604,12 @@ def wait_for_getheaders(self, timeout=60):
         assert wait_until(test_function, timeout=timeout)
 
     def wait_for_inv(self, expected_inv, timeout=60):
-        test_function = lambda: self.last_message.get("inv") and self.last_message["inv"] != expected_inv
+        """Waits for an INV message and checks that the first inv object in the message was as expected."""
+        if len(expected_inv) > 1:
+            raise NotImplementedError("wait_for_inv() will only verify the first inv object")
+        test_function = lambda: self.last_message.get("inv") and \
+                                self.last_message["inv"].inv[0].type == expected_inv[0].type and \
+                                self.last_message["inv"].inv[0].hash == expected_inv[0].hash
         assert wait_until(test_function, timeout=timeout)
 
     def wait_for_verack(self, timeout=60):