test: use script_util helpers for creating P2W{PKH,SH} scripts

6 files changed, 58 insertions, 80 deletions

diff --git a/test/functional/feature_segwit.py b/test/functional/feature_segwit.py
index dadccbd73e..9cf46d9d11 100755
--- a/test/functional/feature_segwit.py
+++ b/test/functional/feature_segwit.py
@@ -23,7 +23,6 @@ from test_framework.messages import (
     CTransaction,
     CTxIn,
     CTxOut,
-    sha256,
     tx_from_hex,
 )
 from test_framework.script import (
@@ -35,11 +34,12 @@ from test_framework.script import (
     OP_CHECKSIG,
     OP_DROP,
     OP_TRUE,
-    hash160,
 )
 from test_framework.script_util import (
     key_to_p2pkh_script,
+    key_to_p2wpkh_script,
     script_to_p2sh_script,
+    script_to_p2wsh_script,
 )
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
@@ -428,9 +428,9 @@ class SegWitTest(BitcoinTestFramework):
         # 2N7MGY19ti4KDMSzRfPAssP6Pxyuxoi6jLe is the P2SH(P2PKH) version of mjoE3sSrb8ByYEvgnC3Aox86u1CHnfJA4V
         unsolvable_address_key = hex_str_to_bytes("02341AEC7587A51CDE5279E0630A531AEA2615A9F80B17E8D9376327BAEAA59E3D")
         unsolvablep2pkh = key_to_p2pkh_script(unsolvable_address_key)
-        unsolvablep2wshp2pkh = CScript([OP_0, sha256(unsolvablep2pkh)])
+        unsolvablep2wshp2pkh = script_to_p2wsh_script(unsolvablep2pkh)
         p2shop0 = script_to_p2sh_script(op0)
-        p2wshop1 = CScript([OP_0, sha256(op1)])
+        p2wshop1 = script_to_p2wsh_script(op1)
         unsolvable_after_importaddress.append(unsolvablep2pkh)
         unsolvable_after_importaddress.append(unsolvablep2wshp2pkh)
         unsolvable_after_importaddress.append(op1)  # OP_1 will be imported as script
@@ -450,16 +450,16 @@ class SegWitTest(BitcoinTestFramework):
             if (v['isscript']):
                 bare = hex_str_to_bytes(v['hex'])
                 importlist.append(bare.hex())
-                importlist.append(CScript([OP_0, sha256(bare)]).hex())
+                importlist.append(script_to_p2wsh_script(bare).hex())
             else:
                 pubkey = hex_str_to_bytes(v['pubkey'])
                 p2pk = CScript([pubkey, OP_CHECKSIG])
                 p2pkh = key_to_p2pkh_script(pubkey)
                 importlist.append(p2pk.hex())
                 importlist.append(p2pkh.hex())
-                importlist.append(CScript([OP_0, hash160(pubkey)]).hex())
-                importlist.append(CScript([OP_0, sha256(p2pk)]).hex())
-                importlist.append(CScript([OP_0, sha256(p2pkh)]).hex())
+                importlist.append(key_to_p2wpkh_script(pubkey).hex())
+                importlist.append(script_to_p2wsh_script(p2pk).hex())
+                importlist.append(script_to_p2wsh_script(p2pkh).hex())
 
         importlist.append(unsolvablep2pkh.hex())
         importlist.append(unsolvablep2wshp2pkh.hex())
@@ -614,20 +614,20 @@ class SegWitTest(BitcoinTestFramework):
     def p2sh_address_to_script(self, v):
         bare = CScript(hex_str_to_bytes(v['hex']))
         p2sh = CScript(hex_str_to_bytes(v['scriptPubKey']))
-        p2wsh = CScript([OP_0, sha256(bare)])
+        p2wsh = script_to_p2wsh_script(bare)
         p2sh_p2wsh = script_to_p2sh_script(p2wsh)
         return([bare, p2sh, p2wsh, p2sh_p2wsh])
 
     def p2pkh_address_to_script(self, v):
         pubkey = hex_str_to_bytes(v['pubkey'])
-        p2wpkh = CScript([OP_0, hash160(pubkey)])
+        p2wpkh = key_to_p2wpkh_script(pubkey)
         p2sh_p2wpkh = script_to_p2sh_script(p2wpkh)
         p2pk = CScript([pubkey, OP_CHECKSIG])
         p2pkh = CScript(hex_str_to_bytes(v['scriptPubKey']))
         p2sh_p2pk = script_to_p2sh_script(p2pk)
         p2sh_p2pkh = script_to_p2sh_script(p2pkh)
-        p2wsh_p2pk = CScript([OP_0, sha256(p2pk)])
-        p2wsh_p2pkh = CScript([OP_0, sha256(p2pkh)])
+        p2wsh_p2pk = script_to_p2wsh_script(p2pk)
+        p2wsh_p2pkh = script_to_p2wsh_script(p2pkh)
         p2sh_p2wsh_p2pk = script_to_p2sh_script(p2wsh_p2pk)
         p2sh_p2wsh_p2pkh = script_to_p2sh_script(p2wsh_p2pkh)
         return [p2wpkh, p2sh_p2wpkh, p2pk, p2pkh, p2sh_p2pk, p2sh_p2pkh, p2wsh_p2pk, p2wsh_p2pkh, p2sh_p2wsh_p2pk, p2sh_p2wsh_p2pkh]
diff --git a/test/functional/feature_taproot.py b/test/functional/feature_taproot.py
index fefa369ffe..17be29c7bf 100755
--- a/test/functional/feature_taproot.py
+++ b/test/functional/feature_taproot.py
@@ -76,15 +76,16 @@ from test_framework.script import (
     taproot_construct,
 )
 from test_framework.script_util import (
+    key_to_p2wpkh_script,
     keyhash_to_p2pkh_script,
     script_to_p2sh_script,
+    script_to_p2wsh_script,
 )
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import assert_raises_rpc_error, assert_equal
 from test_framework.key import generate_privkey, compute_xonly_pubkey, sign_schnorr, tweak_add_privkey, ECKey
 from test_framework.address import (
     hash160,
-    sha256,
 )
 from collections import OrderedDict, namedtuple
 from io import BytesIO
@@ -461,13 +462,13 @@ def make_spender(comment, *, tap=None, witv0=False, script=None, pkh=None, p2sh=
             # P2WPKH
             assert script is None
             pubkeyhash = hash160(pkh)
-            spk = CScript([OP_0, pubkeyhash])
+            spk = key_to_p2wpkh_script(pkh)
             conf["scriptcode"] = keyhash_to_p2pkh_script(pubkeyhash)
             conf["script_witv0"] = None
             conf["inputs"] = [getter("sign"), pkh]
         elif script is not None:
             # P2WSH
-            spk = CScript([OP_0, sha256(script)])
+            spk = script_to_p2wsh_script(script)
             conf["scriptcode"] = script
             conf["script_witv0"] = script
         else:
diff --git a/test/functional/p2p_segwit.py b/test/functional/p2p_segwit.py
index 6a1cb8c1b4..ead9d852fe 100755
--- a/test/functional/p2p_segwit.py
+++ b/test/functional/p2p_segwit.py
@@ -41,7 +41,6 @@ from test_framework.messages import (
     ser_vector,
     sha256,
     tx_from_hex,
-    uint256_from_str,
 )
 from test_framework.p2p import (
     P2PInterface,
@@ -74,8 +73,10 @@ from test_framework.script import (
     hash160,
 )
 from test_framework.script_util import (
+    key_to_p2wpkh_script,
     keyhash_to_p2pkh_script,
     script_to_p2sh_script,
+    script_to_p2wsh_script,
 )
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import (
@@ -488,8 +489,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Create two outputs, a p2wsh and p2sh-p2wsh
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
         p2sh_script_pubkey = script_to_p2sh_script(script_pubkey)
 
         value = self.utxo[0].nValue // 3
@@ -625,8 +625,7 @@ class SegWitTest(BitcoinTestFramework):
         V0 segwit inputs may only be mined after activation, but not before."""
 
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
         p2sh_script_pubkey = script_to_p2sh_script(witness_program)
 
         # First prepare a p2sh output (so that spending it will pass standardness)
@@ -654,6 +653,7 @@ class SegWitTest(BitcoinTestFramework):
         test_transaction_acceptance(self.nodes[1], self.std_node, tx, with_witness=True, accepted=True)
 
         # Now create something that looks like a P2PKH output. This won't be spendable.
+        witness_hash = sha256(witness_program)
         script_pubkey = CScript([OP_0, hash160(witness_hash)])
         tx2 = CTransaction()
         # tx was accepted, so we spend the second output.
@@ -732,8 +732,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Prepare the p2sh-wrapped witness output
         witness_program = CScript([OP_DROP, OP_TRUE])
-        witness_hash = sha256(witness_program)
-        p2wsh_pubkey = CScript([OP_0, witness_hash])
+        p2wsh_pubkey = script_to_p2wsh_script(witness_program)
         script_pubkey = script_to_p2sh_script(p2wsh_pubkey)
         script_sig = CScript([p2wsh_pubkey])  # a push of the redeem script
 
@@ -828,8 +827,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Let's construct a witness program
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
         tx.vout.append(CTxOut(self.utxo[0].nValue - 1000, script_pubkey))
         tx.rehash()
 
@@ -942,8 +940,7 @@ class SegWitTest(BitcoinTestFramework):
         NUM_OUTPUTS = 50
 
         witness_program = CScript([OP_2DROP] * NUM_DROPS + [OP_TRUE])
-        witness_hash = uint256_from_str(sha256(witness_program))
-        script_pubkey = CScript([OP_0, ser_uint256(witness_hash)])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         prevout = COutPoint(self.utxo[0].sha256, self.utxo[0].n)
         value = self.utxo[0].nValue
@@ -1045,8 +1042,7 @@ class SegWitTest(BitcoinTestFramework):
         block = self.build_next_block()
 
         witness_program = CScript([OP_DROP, OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         # First try extra witness data on a tx that doesn't require a witness
         tx = CTransaction()
@@ -1118,8 +1114,7 @@ class SegWitTest(BitcoinTestFramework):
         block = self.build_next_block()
 
         witness_program = CScript([OP_DROP, OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         tx = CTransaction()
         tx.vin.append(CTxIn(COutPoint(self.utxo[0].sha256, self.utxo[0].n), b""))
@@ -1157,8 +1152,7 @@ class SegWitTest(BitcoinTestFramework):
         # This program is 19 max pushes (9937 bytes), then 64 more opcode-bytes.
         long_witness_program = CScript([b'a' * MAX_SCRIPT_ELEMENT_SIZE] * 19 + [OP_DROP] * 63 + [OP_TRUE])
         assert len(long_witness_program) == MAX_PROGRAM_LENGTH + 1
-        long_witness_hash = sha256(long_witness_program)
-        long_script_pubkey = CScript([OP_0, long_witness_hash])
+        long_script_pubkey = script_to_p2wsh_script(long_witness_program)
 
         block = self.build_next_block()
 
@@ -1181,8 +1175,7 @@ class SegWitTest(BitcoinTestFramework):
         # Try again with one less byte in the witness program
         witness_program = CScript([b'a' * MAX_SCRIPT_ELEMENT_SIZE] * 19 + [OP_DROP] * 62 + [OP_TRUE])
         assert len(witness_program) == MAX_PROGRAM_LENGTH
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         tx.vout[0] = CTxOut(tx.vout[0].nValue, script_pubkey)
         tx.rehash()
@@ -1201,8 +1194,7 @@ class SegWitTest(BitcoinTestFramework):
         """Test that vin length must match vtxinwit length."""
 
         witness_program = CScript([OP_DROP, OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         # Create a transaction that splits our utxo into many outputs
         tx = CTransaction()
@@ -1309,8 +1301,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Now try to add extra witness data to a valid witness tx.
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
         tx2 = CTransaction()
         tx2.vin.append(CTxIn(COutPoint(tx_hash, 0), b""))
         tx2.vout.append(CTxOut(tx.vout[0].nValue - 1000, script_pubkey))
@@ -1472,8 +1463,7 @@ class SegWitTest(BitcoinTestFramework):
         block = self.build_next_block()
         # Change the output of the block to be a witness output.
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
         block.vtx[0].vout[0].scriptPubKey = script_pubkey
         # This next line will rehash the coinbase and update the merkle
         # root, and solve.
@@ -1520,7 +1510,7 @@ class SegWitTest(BitcoinTestFramework):
         # Test 1: P2WPKH
         # First create a P2WPKH output that uses an uncompressed pubkey
         pubkeyhash = hash160(pubkey)
-        script_pkh = CScript([OP_0, pubkeyhash])
+        script_pkh = key_to_p2wpkh_script(pubkey)
         tx = CTransaction()
         tx.vin.append(CTxIn(COutPoint(utxo.sha256, utxo.n), b""))
         tx.vout.append(CTxOut(utxo.nValue - 1000, script_pkh))
@@ -1534,8 +1524,7 @@ class SegWitTest(BitcoinTestFramework):
         # Now try to spend it. Send it to a P2WSH output, which we'll
         # use in the next test.
         witness_program = CScript([pubkey, CScriptOp(OP_CHECKSIG)])
-        witness_hash = sha256(witness_program)
-        script_wsh = CScript([OP_0, witness_hash])
+        script_wsh = script_to_p2wsh_script(witness_program)
 
         tx2 = CTransaction()
         tx2.vin.append(CTxIn(COutPoint(tx.sha256, 0), b""))
@@ -1613,8 +1602,7 @@ class SegWitTest(BitcoinTestFramework):
         pubkey = key.get_pubkey().get_bytes()
 
         witness_program = CScript([pubkey, CScriptOp(OP_CHECKSIG)])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         # First create a witness output for use in the tests.
         tx = CTransaction()
@@ -1733,7 +1721,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Now test witness version 0 P2PKH transactions
         pubkeyhash = hash160(pubkey)
-        script_pkh = CScript([OP_0, pubkeyhash])
+        script_pkh = key_to_p2wpkh_script(pubkey)
         tx = CTransaction()
         tx.vin.append(CTxIn(COutPoint(temp_utxos[0].sha256, temp_utxos[0].n), b""))
         tx.vout.append(CTxOut(temp_utxos[0].nValue, script_pkh))
@@ -1860,7 +1848,7 @@ class SegWitTest(BitcoinTestFramework):
         # For each script, generate a pair of P2WSH and P2SH-P2WSH output.
         outputvalue = (self.utxo[0].nValue - 1000) // (len(scripts) * 2)
         for i in scripts:
-            p2wsh = CScript([OP_0, sha256(i)])
+            p2wsh = script_to_p2wsh_script(i)
             p2wsh_scripts.append(p2wsh)
             tx.vout.append(CTxOut(outputvalue, p2wsh))
             tx.vout.append(CTxOut(outputvalue, script_to_p2sh_script(p2wsh)))
@@ -1877,13 +1865,13 @@ class SegWitTest(BitcoinTestFramework):
         for i in range(len(scripts)):
             p2wsh_tx = CTransaction()
             p2wsh_tx.vin.append(CTxIn(COutPoint(txid, i * 2)))
-            p2wsh_tx.vout.append(CTxOut(outputvalue - 5000, CScript([OP_0, hash160(hex_str_to_bytes(""))])))
+            p2wsh_tx.vout.append(CTxOut(outputvalue - 5000, CScript([OP_0, hash160(b"")])))
             p2wsh_tx.wit.vtxinwit.append(CTxInWitness())
             p2wsh_tx.rehash()
             p2wsh_txs.append(p2wsh_tx)
             p2sh_tx = CTransaction()
             p2sh_tx.vin.append(CTxIn(COutPoint(txid, i * 2 + 1), CScript([p2wsh_scripts[i]])))
-            p2sh_tx.vout.append(CTxOut(outputvalue - 5000, CScript([OP_0, hash160(hex_str_to_bytes(""))])))
+            p2sh_tx.vout.append(CTxOut(outputvalue - 5000, CScript([OP_0, hash160(b"")])))
             p2sh_tx.wit.vtxinwit.append(CTxInWitness())
             p2sh_tx.rehash()
             p2sh_txs.append(p2sh_tx)
@@ -1978,8 +1966,7 @@ class SegWitTest(BitcoinTestFramework):
 
         # Keep this under MAX_OPS_PER_SCRIPT (201)
         witness_program = CScript([OP_TRUE, OP_IF, OP_TRUE, OP_ELSE] + [OP_CHECKMULTISIG] * 5 + [OP_CHECKSIG] * 193 + [OP_ENDIF])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         sigops_per_script = 20 * 5 + 193 * 1
         # We'll produce 2 extra outputs, one with a program that would take us
@@ -1995,14 +1982,12 @@ class SegWitTest(BitcoinTestFramework):
         # N(=MAX_SIGOP_COST//sigops_per_script) outputs of our transaction,
         # would push us just over the block sigop limit.
         witness_program_toomany = CScript([OP_TRUE, OP_IF, OP_TRUE, OP_ELSE] + [OP_CHECKSIG] * (extra_sigops_available + 1) + [OP_ENDIF])
-        witness_hash_toomany = sha256(witness_program_toomany)
-        script_pubkey_toomany = CScript([OP_0, witness_hash_toomany])
+        script_pubkey_toomany = script_to_p2wsh_script(witness_program_toomany)
 
         # If we spend this script instead, we would exactly reach our sigop
         # limit (for witness sigops).
         witness_program_justright = CScript([OP_TRUE, OP_IF, OP_TRUE, OP_ELSE] + [OP_CHECKSIG] * (extra_sigops_available) + [OP_ENDIF])
-        witness_hash_justright = sha256(witness_program_justright)
-        script_pubkey_justright = CScript([OP_0, witness_hash_justright])
+        script_pubkey_justright = script_to_p2wsh_script(witness_program_justright)
 
         # First split our available utxo into a bunch of outputs
         split_value = self.utxo[0].nValue // outputs
@@ -2135,8 +2120,7 @@ class SegWitTest(BitcoinTestFramework):
         # Create a Segwit output from the latest UTXO
         # and announce it to the network
         witness_program = CScript([OP_TRUE])
-        witness_hash = sha256(witness_program)
-        script_pubkey = CScript([OP_0, witness_hash])
+        script_pubkey = script_to_p2wsh_script(witness_program)
 
         tx = CTransaction()
         tx.vin.append(CTxIn(COutPoint(self.utxo[0].sha256, self.utxo[0].n), b""))
diff --git a/test/functional/rpc_signrawtransaction.py b/test/functional/rpc_signrawtransaction.py
index ef5d08e7b9..f3627d1e37 100755
--- a/test/functional/rpc_signrawtransaction.py
+++ b/test/functional/rpc_signrawtransaction.py
@@ -6,7 +6,6 @@
 
 from test_framework.blocktools import COINBASE_MATURITY
 from test_framework.address import (
-    check_script,
     script_to_p2sh,
     script_to_p2wsh,
 )
@@ -20,12 +19,10 @@ from test_framework.util import (
 )
 from test_framework.messages import (
     CTxInWitness,
-    sha256,
     tx_from_hex,
 )
 from test_framework.script import (
     CScript,
-    OP_0,
     OP_CHECKLOCKTIMEVERIFY,
     OP_CHECKSIG,
     OP_CHECKSEQUENCEVERIFY,
@@ -233,7 +230,7 @@ class SignRawTransactionsTest(BitcoinTestFramework):
             'P2PKH': key_to_p2pkh_script(embedded_pubkey).hex(),
             'P2PK': CScript([hex_str_to_bytes(embedded_pubkey), OP_CHECKSIG]).hex()
         }.get(tx_type, "Invalid tx_type")
-        redeem_script = CScript([OP_0, sha256(check_script(witness_script))]).hex()
+        redeem_script = script_to_p2wsh_script(witness_script).hex()
         addr = script_to_p2sh(redeem_script)
         script_pub_key = self.nodes[1].validateaddress(addr)['scriptPubKey']
         # Fund that address
diff --git a/test/functional/test_framework/blocktools.py b/test/functional/test_framework/blocktools.py
index bc43438810..833a215993 100644
--- a/test/functional/test_framework/blocktools.py
+++ b/test/functional/test_framework/blocktools.py
@@ -26,7 +26,6 @@ from .messages import (
     hash256,
     hex_str_to_bytes,
     ser_uint256,
-    sha256,
     tx_from_hex,
     uint256_from_str,
 )
@@ -34,13 +33,15 @@ from .script import (
     CScript,
     CScriptNum,
     CScriptOp,
-    OP_0,
     OP_1,
     OP_CHECKMULTISIG,
     OP_CHECKSIG,
     OP_RETURN,
     OP_TRUE,
-    hash160,
+)
+from .script_util import (
+    key_to_p2wpkh_script,
+    script_to_p2wsh_script,
 )
 from .util import assert_equal
 
@@ -206,13 +207,11 @@ def witness_script(use_p2wsh, pubkey):
     scriptPubKey."""
     if not use_p2wsh:
         # P2WPKH instead
-        pubkeyhash = hash160(hex_str_to_bytes(pubkey))
-        pkscript = CScript([OP_0, pubkeyhash])
+        pkscript = key_to_p2wpkh_script(pubkey)
     else:
         # 1-of-1 multisig
         witness_program = CScript([OP_1, hex_str_to_bytes(pubkey), OP_1, OP_CHECKMULTISIG])
-        scripthash = sha256(witness_program)
-        pkscript = CScript([OP_0, scripthash])
+        pkscript = script_to_p2wsh_script(witness_program)
     return pkscript.hex()
 
 def create_witness_tx(node, use_p2wsh, utxo, pubkey, encode_p2sh, amount):
diff --git a/test/functional/test_framework/wallet_util.py b/test/functional/test_framework/wallet_util.py
index 76f1a1edfc..acbc040741 100755
--- a/test/functional/test_framework/wallet_util.py
+++ b/test/functional/test_framework/wallet_util.py
@@ -17,16 +17,15 @@ from test_framework.address import (
 from test_framework.key import ECKey
 from test_framework.script import (
     CScript,
-    OP_0,
     OP_2,
     OP_3,
     OP_CHECKMULTISIG,
-    hash160,
-    sha256,
 )
 from test_framework.script_util import (
     key_to_p2pkh_script,
+    key_to_p2wpkh_script,
     script_to_p2sh_script,
+    script_to_p2wsh_script,
 )
 from test_framework.util import hex_str_to_bytes
 
@@ -56,15 +55,14 @@ def get_key(node):
     Returns a named tuple of privkey, pubkey and all address and scripts."""
     addr = node.getnewaddress()
     pubkey = node.getaddressinfo(addr)['pubkey']
-    pkh = hash160(hex_str_to_bytes(pubkey))
     return Key(privkey=node.dumpprivkey(addr),
                pubkey=pubkey,
                p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
                p2pkh_addr=key_to_p2pkh(pubkey),
-               p2wpkh_script=CScript([OP_0, pkh]).hex(),
+               p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
                p2wpkh_addr=key_to_p2wpkh(pubkey),
-               p2sh_p2wpkh_script=script_to_p2sh_script(CScript([OP_0, pkh])).hex(),
-               p2sh_p2wpkh_redeem_script=CScript([OP_0, pkh]).hex(),
+               p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
+               p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
                p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
 
 def get_generate_key():
@@ -75,15 +73,14 @@ def get_generate_key():
     eckey.generate()
     privkey = bytes_to_wif(eckey.get_bytes())
     pubkey = eckey.get_pubkey().get_bytes().hex()
-    pkh = hash160(hex_str_to_bytes(pubkey))
     return Key(privkey=privkey,
                pubkey=pubkey,
                p2pkh_script=key_to_p2pkh_script(pubkey).hex(),
                p2pkh_addr=key_to_p2pkh(pubkey),
-               p2wpkh_script=CScript([OP_0, pkh]).hex(),
+               p2wpkh_script=key_to_p2wpkh_script(pubkey).hex(),
                p2wpkh_addr=key_to_p2wpkh(pubkey),
-               p2sh_p2wpkh_script=script_to_p2sh_script(CScript([OP_0, pkh])).hex(),
-               p2sh_p2wpkh_redeem_script=CScript([OP_0, pkh]).hex(),
+               p2sh_p2wpkh_script=script_to_p2sh_script(key_to_p2wpkh_script(pubkey)).hex(),
+               p2sh_p2wpkh_redeem_script=key_to_p2wpkh_script(pubkey).hex(),
                p2sh_p2wpkh_addr=key_to_p2sh_p2wpkh(pubkey))
 
 def get_multisig(node):
@@ -97,7 +94,7 @@ def get_multisig(node):
         addrs.append(addr['address'])
         pubkeys.append(addr['pubkey'])
     script_code = CScript([OP_2] + [hex_str_to_bytes(pubkey) for pubkey in pubkeys] + [OP_3, OP_CHECKMULTISIG])
-    witness_script = CScript([OP_0, sha256(script_code)])
+    witness_script = script_to_p2wsh_script(script_code)
     return Multisig(privkeys=[node.dumpprivkey(addr) for addr in addrs],
                     pubkeys=pubkeys,
                     p2sh_script=script_to_p2sh_script(script_code).hex(),