From 3451203b5c67c234d168a409d69ff9623573dae3 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Mon, 3 Oct 2016 18:30:52 -0400
Subject: [PATCH] [qa] Respond to getheaders and do not assume a getdata on inv

---
 qa/rpc-tests/p2p-compactblocks.py |  7 +++++++
 qa/rpc-tests/p2p-segwit.py        | 14 ++++++++++++--
 qa/rpc-tests/sendheaders.py       | 11 +++++------
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/qa/rpc-tests/p2p-compactblocks.py b/qa/rpc-tests/p2p-compactblocks.py
index 6b5d4771310..722ec0f62e7 100755
--- a/qa/rpc-tests/p2p-compactblocks.py
+++ b/qa/rpc-tests/p2p-compactblocks.py
@@ -27,6 +27,7 @@ class TestNode(SingleNodeConnCB):
         self.last_cmpctblock = None
         self.block_announced = False
         self.last_getdata = None
+        self.last_getheaders = None
         self.last_getblocktxn = None
         self.last_block = None
         self.last_blocktxn = None
@@ -64,6 +65,9 @@ class TestNode(SingleNodeConnCB):
     def on_getdata(self, conn, message):
         self.last_getdata = message
 
+    def on_getheaders(self, conn, message):
+        self.last_getheaders = message
+
     def on_getblocktxn(self, conn, message):
         self.last_getblocktxn = message
 
@@ -393,6 +397,9 @@ class CompactBlocksTest(BitcoinTestFramework):
 
             if announce == "inv":
                 test_node.send_message(msg_inv([CInv(2, block.sha256)]))
+                success = wait_until(lambda: test_node.last_getheaders is not None, timeout=30)
+                assert(success)
+                test_node.send_header_for_blocks([block])
             else:
                 test_node.send_header_for_blocks([block])
             success = wait_until(lambda: test_node.last_getdata is not None, timeout=30)
diff --git a/qa/rpc-tests/p2p-segwit.py b/qa/rpc-tests/p2p-segwit.py
index 09ab1b80fcb..6ecd84c3f08 100755
--- a/qa/rpc-tests/p2p-segwit.py
+++ b/qa/rpc-tests/p2p-segwit.py
@@ -64,6 +64,9 @@ class TestNode(NodeConnCB):
             self.getdataset.add(inv.hash)
         self.last_getdata = message
 
+    def on_getheaders(self, conn, message):
+        self.last_getheaders = message
+
     def on_pong(self, conn, message):
         self.last_pong = message
 
@@ -97,6 +100,10 @@ class TestNode(NodeConnCB):
         test_function = lambda: self.last_getdata != None
         self.sync(test_function, timeout)
 
+    def wait_for_getheaders(self, timeout=60):
+        test_function = lambda: self.last_getheaders != None
+        self.sync(test_function, timeout)
+
     def wait_for_inv(self, expected_inv, timeout=60):
         test_function = lambda: self.last_inv != expected_inv
         self.sync(test_function, timeout)
@@ -111,12 +118,15 @@ class TestNode(NodeConnCB):
     def announce_block_and_wait_for_getdata(self, block, use_header, timeout=60):
         with mininode_lock:
             self.last_getdata = None
+            self.last_getheaders = None
+        msg = msg_headers()
+        msg.headers = [ CBlockHeader(block) ]
         if use_header:
-            msg = msg_headers()
-            msg.headers = [ CBlockHeader(block) ]
             self.send_message(msg)
         else:
             self.send_message(msg_inv(inv=[CInv(2, block.sha256)]))
+            self.wait_for_getheaders()
+            self.send_message(msg)
         self.wait_for_getdata()
         return
 
diff --git a/qa/rpc-tests/sendheaders.py b/qa/rpc-tests/sendheaders.py
index 81b2442e6ae..37b98c576e8 100755
--- a/qa/rpc-tests/sendheaders.py
+++ b/qa/rpc-tests/sendheaders.py
@@ -348,14 +348,13 @@ class SendHeadersTest(BitcoinTestFramework):
                 if j == 0:
                     # Announce via inv
                     test_node.send_block_inv(tip)
-                    test_node.wait_for_getdata([tip], timeout=5)
+                    test_node.wait_for_getheaders(timeout=5)
+                    # Should have received a getheaders now
+                    test_node.send_header_for_blocks(blocks)
                     # Test that duplicate inv's won't result in duplicate
                     # getdata requests, or duplicate headers announcements
-                    inv_node.send_block_inv(tip)
-                    # Should have received a getheaders as well!
-                    test_node.send_header_for_blocks(blocks)
-                    test_node.wait_for_getdata([x.sha256 for x in blocks[0:-1]], timeout=5)
-                    [ inv_node.send_block_inv(x.sha256) for x in blocks[0:-1] ]
+                    [ inv_node.send_block_inv(x.sha256) for x in blocks ]
+                    test_node.wait_for_getdata([x.sha256 for x in blocks], timeout=5)
                     inv_node.sync_with_ping()
                 else:
                     # Announce via headers