[qa] Remove hardcoded "4 nodes" from test_framework

2 files changed, 33 insertions, 25 deletions

diff --git a/qa/rpc-tests/test_framework/test_framework.py b/qa/rpc-tests/test_framework/test_framework.py
index 3480de6c6e..b9b803429d 100755
--- a/qa/rpc-tests/test_framework/test_framework.py
+++ b/qa/rpc-tests/test_framework/test_framework.py
@@ -15,7 +15,6 @@ import traceback
 
 from .util import (
     initialize_chain,
-    assert_equal,
     start_nodes,
     connect_nodes_bi,
     sync_blocks,
@@ -32,21 +31,26 @@ from .authproxy import JSONRPCException
 
 class BitcoinTestFramework(object):
 
-    # These may be over-ridden by subclasses:
+    def __init__(self):
+        self.num_nodes = 4
+        self.setup_clean_chain = False
+        self.nodes = None
+
     def run_test(self):
-        for node in self.nodes:
-            assert_equal(node.getblockcount(), 200)
-            assert_equal(node.getbalance(), 25*50)
+        raise NotImplementedError
 
     def add_options(self, parser):
         pass
 
     def setup_chain(self):
         print("Initializing test directory "+self.options.tmpdir)
-        initialize_chain(self.options.tmpdir)
+        if self.setup_clean_chain:
+            initialize_chain_clean(self.options.tmpdir, self.num_nodes)
+        else:
+            initialize_chain(self.options.tmpdir, self.num_nodes)
 
     def setup_nodes(self):
-        return start_nodes(4, self.options.tmpdir)
+        return start_nodes(self.num_nodes, self.options.tmpdir)
 
     def setup_network(self, split = False):
         self.nodes = self.setup_nodes()
@@ -181,9 +185,10 @@ class BitcoinTestFramework(object):
 
 class ComparisonTestFramework(BitcoinTestFramework):
 
-    # Can override the num_nodes variable to indicate how many nodes to run.
     def __init__(self):
+        super().__init__()
         self.num_nodes = 2
+        self.setup_clean_chain = True
 
     def add_options(self, parser):
         parser.add_option("--testbinary", dest="testbinary",
@@ -193,10 +198,6 @@ class ComparisonTestFramework(BitcoinTestFramework):
                           default=os.getenv("BITCOIND", "bitcoind"),
                           help="bitcoind binary to use for reference nodes (if any)")
 
-    def setup_chain(self):
-        print("Initializing test directory "+self.options.tmpdir)
-        initialize_chain_clean(self.options.tmpdir, self.num_nodes)
-
     def setup_network(self):
         self.nodes = start_nodes(
             self.num_nodes, self.options.tmpdir,
diff --git a/qa/rpc-tests/test_framework/util.py b/qa/rpc-tests/test_framework/util.py
index 6784177aaa..0e522e78d4 100644
--- a/qa/rpc-tests/test_framework/util.py
+++ b/qa/rpc-tests/test_framework/util.py
@@ -182,24 +182,28 @@ def wait_for_bitcoind_start(process, url, i):
                 raise # unkown JSON RPC exception
         time.sleep(0.25)
 
-def initialize_chain(test_dir):
+def initialize_chain(test_dir, num_nodes):
     """
-    Create (or copy from cache) a 200-block-long chain and
-    4 wallets.
+    Create a cache of a 200-block-long chain (with wallet) for MAX_NODES
+    Afterward, create num_nodes copies from the cache
     """
 
-    if (not os.path.isdir(os.path.join("cache","node0"))
-        or not os.path.isdir(os.path.join("cache","node1"))
-        or not os.path.isdir(os.path.join("cache","node2"))
-        or not os.path.isdir(os.path.join("cache","node3"))):
+    assert num_nodes <= MAX_NODES
+    create_cache = False
+    for i in range(MAX_NODES):
+        if not os.path.isdir(os.path.join('cache', 'node'+str(i))):
+            create_cache = True
+            break
+
+    if create_cache:
 
         #find and delete old cache directories if any exist
-        for i in range(4):
+        for i in range(MAX_NODES):
             if os.path.isdir(os.path.join("cache","node"+str(i))):
                 shutil.rmtree(os.path.join("cache","node"+str(i)))
 
         # Create cache directories, run bitcoinds:
-        for i in range(4):
+        for i in range(MAX_NODES):
             datadir=initialize_datadir("cache", i)
             args = [ os.getenv("BITCOIND", "bitcoind"), "-server", "-keypool=1", "-datadir="+datadir, "-discover=0" ]
             if i > 0:
@@ -212,15 +216,18 @@ def initialize_chain(test_dir):
                 print("initialize_chain: RPC succesfully started")
 
         rpcs = []
-        for i in range(4):
+        for i in range(MAX_NODES):
             try:
                 rpcs.append(get_rpc_proxy(rpc_url(i), i))
             except:
                 sys.stderr.write("Error connecting to "+url+"\n")
                 sys.exit(1)
 
-        # Create a 200-block-long chain; each of the 4 nodes
+        # Create a 200-block-long chain; each of the 4 first nodes
         # gets 25 mature blocks and 25 immature.
+        # Note: To preserve compatibility with older versions of
+        # initialize_chain, only 4 nodes will generate coins.
+        #
         # blocks are created with timestamps 10 minutes apart
         # starting from 2010 minutes in the past
         enable_mocktime()
@@ -238,13 +245,13 @@ def initialize_chain(test_dir):
         stop_nodes(rpcs)
         wait_bitcoinds()
         disable_mocktime()
-        for i in range(4):
+        for i in range(MAX_NODES):
             os.remove(log_filename("cache", i, "debug.log"))
             os.remove(log_filename("cache", i, "db.log"))
             os.remove(log_filename("cache", i, "peers.dat"))
             os.remove(log_filename("cache", i, "fee_estimates.dat"))
 
-    for i in range(4):
+    for i in range(num_nodes):
         from_dir = os.path.join("cache", "node"+str(i))
         to_dir = os.path.join(test_dir,  "node"+str(i))
         shutil.copytree(from_dir, to_dir)