From 5ebd5f9e1528e8c185e83981007dd6198c7f73dd Mon Sep 17 00:00:00 2001
From: John Newbery <john@johnnewbery.com>
Date: Wed, 7 Jun 2017 16:14:20 -0400
Subject: [PATCH] [tests] tidy up zmq_test.py

---
 test/functional/test_runner.py |  2 +-
 test/functional/zmq_test.py    | 49 ++++++++++++++++------------------
 2 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/test/functional/test_runner.py b/test/functional/test_runner.py
index 4702f2d773..8ccb172d9b 100755
--- a/test/functional/test_runner.py
+++ b/test/functional/test_runner.py
@@ -78,7 +78,7 @@ BASE_SCRIPTS= [
     'rawtransactions.py',
     'reindex.py',
     # vv Tests less than 30s vv
-    "zmq_test.py",
+    'zmq_test.py',
     'mempool_resurrect_test.py',
     'txn_doublespend.py --mineblock',
     'txn_clone.py',
diff --git a/test/functional/zmq_test.py b/test/functional/zmq_test.py
index ca89e200ed..fd7cbc05da 100755
--- a/test/functional/zmq_test.py
+++ b/test/functional/zmq_test.py
@@ -8,15 +8,15 @@ import os
 import struct
 
 from test_framework.test_framework import BitcoinTestFramework, SkipTest
-from test_framework.util import *
+from test_framework.util import (assert_equal,
+                                 bytes_to_hex_str,
+                                 )
 
 class ZMQTest (BitcoinTestFramework):
 
     def __init__(self):
         super().__init__()
-        self.num_nodes = 4
-
-    port = 28332
+        self.num_nodes = 2
 
     def setup_nodes(self):
         # Try to import python3-zmq. Skip this test if the import fails.
@@ -38,57 +38,55 @@ class ZMQTest (BitcoinTestFramework):
         self.zmqSubSocket = self.zmqContext.socket(zmq.SUB)
         self.zmqSubSocket.setsockopt(zmq.SUBSCRIBE, b"hashblock")
         self.zmqSubSocket.setsockopt(zmq.SUBSCRIBE, b"hashtx")
-        self.zmqSubSocket.connect("tcp://127.0.0.1:%i" % self.port)
-        self.nodes = self.start_nodes(self.num_nodes, self.options.tmpdir, extra_args=[
-            ['-zmqpubhashtx=tcp://127.0.0.1:'+str(self.port), '-zmqpubhashblock=tcp://127.0.0.1:'+str(self.port)],
-            [],
-            [],
-            []
-            ])
+        ip_address = "tcp://127.0.0.1:28332"
+        self.zmqSubSocket.connect(ip_address)
+        extra_args = [['-zmqpubhashtx=%s' % ip_address, '-zmqpubhashblock=%s' % ip_address], []]
+        self.nodes = self.start_nodes(self.num_nodes, self.options.tmpdir, extra_args)
 
     def run_test(self):
-        self.sync_all()
-
         genhashes = self.nodes[0].generate(1)
         self.sync_all()
 
-        self.log.info("listen...")
+        self.log.info("Wait for tx")
         msg = self.zmqSubSocket.recv_multipart()
         topic = msg[0]
         assert_equal(topic, b"hashtx")
         body = msg[1]
         msgSequence = struct.unpack('<I', msg[-1])[-1]
-        assert_equal(msgSequence, 0) #must be sequence 0 on hashtx
+        assert_equal(msgSequence, 0)  # must be sequence 0 on hashtx
 
+        self.log.info("Wait for block")
         msg = self.zmqSubSocket.recv_multipart()
         topic = msg[0]
         body = msg[1]
         msgSequence = struct.unpack('<I', msg[-1])[-1]
-        assert_equal(msgSequence, 0) #must be sequence 0 on hashblock
+        assert_equal(msgSequence, 0)  # must be sequence 0 on hashblock
         blkhash = bytes_to_hex_str(body)
 
-        assert_equal(genhashes[0], blkhash) #blockhash from generate must be equal to the hash received over zmq
+        assert_equal(genhashes[0], blkhash)  # blockhash from generate must be equal to the hash received over zmq
 
+        self.log.info("Generate 10 blocks (and 10 coinbase txes)")
         n = 10
         genhashes = self.nodes[1].generate(n)
         self.sync_all()
 
         zmqHashes = []
         blockcount = 0
-        for x in range(0,n*2):
+        for x in range(n * 2):
             msg = self.zmqSubSocket.recv_multipart()
             topic = msg[0]
             body = msg[1]
             if topic == b"hashblock":
                 zmqHashes.append(bytes_to_hex_str(body))
                 msgSequence = struct.unpack('<I', msg[-1])[-1]
-                assert_equal(msgSequence, blockcount+1)
+                assert_equal(msgSequence, blockcount + 1)
                 blockcount += 1
 
-        for x in range(0,n):
-            assert_equal(genhashes[x], zmqHashes[x]) #blockhash from generate must be equal to the hash received over zmq
+        for x in range(n):
+            assert_equal(genhashes[x], zmqHashes[x])  # blockhash from generate must be equal to the hash received over zmq
 
-        #test tx from a second node
+        self.log.info("Wait for tx from second node")
+        # test tx from a second node
         hashRPC = self.nodes[1].sendtoaddress(self.nodes[0].getnewaddress(), 1.0)
         self.sync_all()
 
@@ -100,10 +98,9 @@ class ZMQTest (BitcoinTestFramework):
         if topic == b"hashtx":
             hashZMQ = bytes_to_hex_str(body)
             msgSequence = struct.unpack('<I', msg[-1])[-1]
-            assert_equal(msgSequence, blockcount+1)
-
-        assert_equal(hashRPC, hashZMQ) #blockhash from generate must be equal to the hash received over zmq
+            assert_equal(msgSequence, blockcount + 1)
 
+        assert_equal(hashRPC, hashZMQ)  # txid from sendtoaddress must be equal to the hash received over zmq
 
 if __name__ == '__main__':
-    ZMQTest ().main ()
+    ZMQTest().main()