From 9674be8f96c9d5930fb4ad7d498be5b4279f2f3a Mon Sep 17 00:00:00 2001
From: Alexander Block <ablock84@gmail.com>
Date: Wed, 18 Jul 2018 12:06:56 +0200
Subject: [PATCH] Refactor TestChain100Setup to allow different test chain
 setups

Allows testing of features which depend on BIP9 deployments
---
 src/test/test_dash.cpp | 37 +++++++++++++++++++++++++++++--------
 src/test/test_dash.h   | 26 ++++++++++++++++++--------
 2 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/src/test/test_dash.cpp b/src/test/test_dash.cpp
index 5eb785557..01c73f7d9 100644
--- a/src/test/test_dash.cpp
+++ b/src/test/test_dash.cpp
@@ -100,12 +100,12 @@ TestingSetup::~TestingSetup()
         boost::filesystem::remove_all(pathTemp);
 }
 
-TestChain100Setup::TestChain100Setup() : TestingSetup(CBaseChainParams::REGTEST)
+TestChainSetup::TestChainSetup(int blockCount) : TestingSetup(CBaseChainParams::REGTEST)
 {
     // Generate a 100-block chain:
     coinbaseKey.MakeNewKey(true);
-    CScript scriptPubKey = CScript() <<  ToByteVector(coinbaseKey.GetPubKey()) << OP_CHECKSIG;
-    for (int i = 0; i < COINBASE_MATURITY; i++)
+    CScript scriptPubKey = CScript() << ToByteVector(coinbaseKey.GetPubKey()) << OP_CHECKSIG;
+    for (int i = 0; i < blockCount; i++)
     {
         std::vector<CMutableTransaction> noTxns;
         CBlock b = CreateAndProcessBlock(noTxns, scriptPubKey);
@@ -118,7 +118,25 @@ TestChain100Setup::TestChain100Setup() : TestingSetup(CBaseChainParams::REGTEST)
 // scriptPubKey, and try to add it to the current chain.
 //
 CBlock
-TestChain100Setup::CreateAndProcessBlock(const std::vector<CMutableTransaction>& txns, const CScript& scriptPubKey)
+TestChainSetup::CreateAndProcessBlock(const std::vector<CMutableTransaction>& txns, const CScript& scriptPubKey)
+{
+    const CChainParams& chainparams = Params();
+    auto block = CreateBlock(txns, scriptPubKey);
+
+    std::shared_ptr<const CBlock> shared_pblock = std::make_shared<const CBlock>(block);
+    ProcessNewBlock(chainparams, shared_pblock, true, NULL);
+
+    CBlock result = block;
+    return result;
+}
+
+CBlock TestChainSetup::CreateAndProcessBlock(const std::vector<CMutableTransaction>& txns, const CKey& scriptKey)
+{
+    CScript scriptPubKey = CScript() <<  ToByteVector(coinbaseKey.GetPubKey()) << OP_CHECKSIG;
+    return CreateAndProcessBlock(txns, scriptPubKey);
+}
+
+CBlock TestChainSetup::CreateBlock(const std::vector<CMutableTransaction>& txns, const CScript& scriptPubKey)
 {
     const CChainParams& chainparams = Params();
     std::unique_ptr<CBlockTemplate> pblocktemplate = BlockAssembler(chainparams).CreateNewBlock(scriptPubKey);
@@ -134,14 +152,17 @@ TestChain100Setup::CreateAndProcessBlock(const std::vector<CMutableTransaction>&
 
     while (!CheckProofOfWork(block.GetHash(), block.nBits, chainparams.GetConsensus())) ++block.nNonce;
 
-    std::shared_ptr<const CBlock> shared_pblock = std::make_shared<const CBlock>(block);
-    ProcessNewBlock(chainparams, shared_pblock, true, NULL);
-
     CBlock result = block;
     return result;
 }
 
-TestChain100Setup::~TestChain100Setup()
+CBlock TestChainSetup::CreateBlock(const std::vector<CMutableTransaction>& txns, const CKey& scriptKey)
+{
+    CScript scriptPubKey = CScript() <<  ToByteVector(coinbaseKey.GetPubKey()) << OP_CHECKSIG;
+    return CreateBlock(txns, scriptPubKey);
+}
+
+TestChainSetup::~TestChainSetup()
 {
 }
 
diff --git a/src/test/test_dash.h b/src/test/test_dash.h
index 0bb1717a8..1a42fbff5 100644
--- a/src/test/test_dash.h
+++ b/src/test/test_dash.h
@@ -43,24 +43,34 @@ class CBlock;
 struct CMutableTransaction;
 class CScript;
 
-//
-// Testing fixture that pre-creates a
-// 100-block REGTEST-mode block chain
-//
-struct TestChain100Setup : public TestingSetup {
-    TestChain100Setup();
+struct TestChainSetup : public TestingSetup
+{
+    TestChainSetup(int blockCount);
+    ~TestChainSetup();
 
     // Create a new block with just given transactions, coinbase paying to
     // scriptPubKey, and try to add it to the current chain.
     CBlock CreateAndProcessBlock(const std::vector<CMutableTransaction>& txns,
                                  const CScript& scriptPubKey);
-
-    ~TestChain100Setup();
+    CBlock CreateAndProcessBlock(const std::vector<CMutableTransaction>& txns,
+                                 const CKey& scriptKey);
+    CBlock CreateBlock(const std::vector<CMutableTransaction>& txns,
+                       const CScript& scriptPubKey);
+    CBlock CreateBlock(const std::vector<CMutableTransaction>& txns,
+                       const CKey& scriptKey);
 
     std::vector<CTransaction> coinbaseTxns; // For convenience, coinbase transactions
     CKey coinbaseKey; // private/public key needed to spend coinbase transactions
 };
 
+//
+// Testing fixture that pre-creates a
+// 100-block REGTEST-mode block chain
+//
+struct TestChain100Setup : public TestChainSetup {
+    TestChain100Setup() : TestChainSetup(100) {}
+};
+
 class CTxMemPoolEntry;
 class CTxMemPool;