From c39a6b9ec8ea6599639435378e0078218a4155fc Mon Sep 17 00:00:00 2001
From: Jimmy Song <jaejoon@gmail.com>
Date: Mon, 17 Apr 2017 09:23:44 -0700
Subject: [PATCH] Tests: Refactor to create witness script creation function

* Refactor blocktools.py so that witness script creation is its own function
* Changed p2p-segwit to use new function
---
 test/functional/p2p-segwit.py                | 13 ++++---------
 test/functional/test_framework/blocktools.py | 11 ++++++++---
 test/functional/test_framework/mininode.py   |  3 ++-
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/test/functional/p2p-segwit.py b/test/functional/p2p-segwit.py
index cd7b788eb4..3dd78c8b7d 100755
--- a/test/functional/p2p-segwit.py
+++ b/test/functional/p2p-segwit.py
@@ -8,7 +8,7 @@ from test_framework.mininode import *
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import *
 from test_framework.script import *
-from test_framework.blocktools import create_block, create_coinbase, add_witness_commitment, WITNESS_COMMITMENT_HEADER
+from test_framework.blocktools import create_block, create_coinbase, add_witness_commitment, get_witness_script, WITNESS_COMMITMENT_HEADER
 from test_framework.key import CECKey, CPubKey
 import time
 import random
@@ -1721,15 +1721,10 @@ class SegWitTest(BitcoinTestFramework):
                 assert('default_witness_commitment' in gbt_results)
                 witness_commitment = gbt_results['default_witness_commitment']
 
-                # TODO: this duplicates some code from blocktools.py, would be nice
-                # to refactor.
                 # Check that default_witness_commitment is present.
-                block = CBlock()
-                witness_root = block.get_merkle_root([ser_uint256(0), ser_uint256(txid)])
-                check_commitment = uint256_from_str(hash256(ser_uint256(witness_root)+ser_uint256(0)))
-                from test_framework.blocktools import WITNESS_COMMITMENT_HEADER
-                output_data = WITNESS_COMMITMENT_HEADER + ser_uint256(check_commitment)
-                script = CScript([OP_RETURN, output_data])
+                witness_root = CBlock.get_merkle_root([ser_uint256(0),
+                                                       ser_uint256(txid)])
+                script = get_witness_script(witness_root, 0)
                 assert_equal(witness_commitment, bytes_to_hex_str(script))
 
         # undo mocktime
diff --git a/test/functional/test_framework/blocktools.py b/test/functional/test_framework/blocktools.py
index 2c9a0857df..5dcf516dc6 100644
--- a/test/functional/test_framework/blocktools.py
+++ b/test/functional/test_framework/blocktools.py
@@ -25,6 +25,13 @@ def create_block(hashprev, coinbase, nTime=None):
 # From BIP141
 WITNESS_COMMITMENT_HEADER = b"\xaa\x21\xa9\xed"
 
+
+def get_witness_script(witness_root, witness_nonce):
+    witness_commitment = uint256_from_str(hash256(ser_uint256(witness_root)+ser_uint256(witness_nonce)))
+    output_data = WITNESS_COMMITMENT_HEADER + ser_uint256(witness_commitment)
+    return CScript([OP_RETURN, output_data])
+
+
 # According to BIP141, blocks with witness rules active must commit to the
 # hash of all in-block transactions including witness.
 def add_witness_commitment(block, nonce=0):
@@ -32,14 +39,12 @@ def add_witness_commitment(block, nonce=0):
     # transactions, with witnesses.
     witness_nonce = nonce
     witness_root = block.calc_witness_merkle_root()
-    witness_commitment = uint256_from_str(hash256(ser_uint256(witness_root)+ser_uint256(witness_nonce)))
     # witness_nonce should go to coinbase witness.
     block.vtx[0].wit.vtxinwit = [CTxInWitness()]
     block.vtx[0].wit.vtxinwit[0].scriptWitness.stack = [ser_uint256(witness_nonce)]
 
     # witness commitment is the last OP_RETURN output in coinbase
-    output_data = WITNESS_COMMITMENT_HEADER + ser_uint256(witness_commitment)
-    block.vtx[0].vout.append(CTxOut(0, CScript([OP_RETURN, output_data])))
+    block.vtx[0].vout.append(CTxOut(0, get_witness_script(witness_root, witness_nonce)))
     block.vtx[0].rehash()
     block.hashMerkleRoot = block.calc_merkle_root()
     block.rehash()
diff --git a/test/functional/test_framework/mininode.py b/test/functional/test_framework/mininode.py
index d57d46f2ff..2383cca58e 100755
--- a/test/functional/test_framework/mininode.py
+++ b/test/functional/test_framework/mininode.py
@@ -610,7 +610,8 @@ class CBlock(CBlockHeader):
         return r
 
     # Calculate the merkle root given a vector of transaction hashes
-    def get_merkle_root(self, hashes):
+    @classmethod
+    def get_merkle_root(cls, hashes):
         while len(hashes) > 1:
             newhashes = []
             for i in range(0, len(hashes), 2):