test: MiniWallet: introduce enum type for output mode

theStack · theStack · commit 6cebac598e5e · 2021-05-24T16:10:49.000+02:00
For the MiniWallet constructor, the two boolean parameters
"raw_script" and "use_p2pk" are replaced by a single parameter of the
newly introduced type MiniWalletMode (derived by enum.Enum), which can
hold the following values:
	- ADDRESS_OP_TRUE
	- RAW_OP_TRUE
	- RAW_P2PK
diff --git a/test/functional/feature_cltv.py b/test/functional/feature_cltv.py
@@ -26,7 +26,10 @@
 )
 from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import assert_equal
-from test_framework.wallet import MiniWallet
+from test_framework.wallet import (
+    MiniWallet,
+    MiniWalletMode,
+)
 
 CLTV_HEIGHT = 1351
 
@@ -97,7 +100,7 @@ def test_cltv_info(self, *, is_active):
 
     def run_test(self):
         peer = self.nodes[0].add_p2p_connection(P2PInterface())
-        wallet = MiniWallet(self.nodes[0], raw_script=True)
+        wallet = MiniWallet(self.nodes[0], mode=MiniWalletMode.RAW_OP_TRUE)
 
         self.test_cltv_info(is_active=False)
 
diff --git a/test/functional/feature_csv_activation.py b/test/functional/feature_csv_activation.py
@@ -55,7 +55,10 @@
     assert_equal,
     softfork_active,
 )
-from test_framework.wallet import MiniWallet
+from test_framework.wallet import (
+    MiniWallet,
+    MiniWalletMode,
+)
 
 TESTING_TX_COUNT = 83  # Number of testing transactions: 1 BIP113 tx, 16 BIP68 txs, 66 BIP112 txs (see comments above)
 COINBASE_BLOCK_COUNT = TESTING_TX_COUNT  # Number of coinbase blocks we need to generate as inputs for our txs
@@ -181,7 +184,7 @@ def send_blocks(self, blocks, success=True, reject_reason=None):
 
     def run_test(self):
         self.helper_peer = self.nodes[0].add_p2p_connection(P2PDataStore())
-        self.miniwallet = MiniWallet(self.nodes[0], use_p2pk=True)
+        self.miniwallet = MiniWallet(self.nodes[0], mode=MiniWalletMode.RAW_P2PK)
 
         self.log.info("Generate blocks in the past for coinbase outputs.")
         long_past_time = int(time.time()) - 600 * 1000  # enough to build up to 1000 blocks 10 minutes apart without worrying about getting into the future
diff --git a/test/functional/test_framework/wallet.py b/test/functional/test_framework/wallet.py
@@ -5,6 +5,7 @@
 """A limited-functionality wallet, which may replace a real wallet in tests"""
 
 from decimal import Decimal
+from enum import Enum
 from test_framework.address import ADDRESS_BCRT1_P2WSH_OP_TRUE
 from test_framework.key import ECKey
 from test_framework.messages import (
@@ -30,22 +31,46 @@
 )
 
 
+class MiniWalletMode(Enum):
+    """Determines the transaction type the MiniWallet is creating and spending.
+
+    For most purposes, the default mode ADDRESS_OP_TRUE should be sufficient;
+    it simply uses a fixed bech32 P2WSH address whose coins are spent with a
+    witness stack of OP_TRUE, i.e. following an anyone-can-spend policy.
+    However, if the transactions need to be modified by the user (e.g. prepending
+    scriptSig for testing opcodes that are activated by a soft-fork), or the txs
+    should contain an actual signature, the raw modes RAW_OP_TRUE and RAW_P2PK
+    can be useful. Summary of modes:
+
+                    |      output       |           |  tx is   | can modify |  needs
+         mode       |    description    |  address  | standard | scriptSig  | signing
+    ----------------+-------------------+-----------+----------+------------+----------
+    ADDRESS_OP_TRUE | anyone-can-spend  |  bech32   |   yes    |    no      |   no
+    RAW_OP_TRUE     | anyone-can-spend  |  - (raw)  |   no     |    yes     |   no
+    RAW_P2PK        | pay-to-public-key |  - (raw)  |   yes    |    yes     |   yes
+    """
+    ADDRESS_OP_TRUE = 1
+    RAW_OP_TRUE = 2
+    RAW_P2PK = 3
+
+
 class MiniWallet:
-    def __init__(self, test_node, *, raw_script=False, use_p2pk=False):
+    def __init__(self, test_node, *, mode=MiniWalletMode.ADDRESS_OP_TRUE):
         self._test_node = test_node
         self._utxos = []
         self._priv_key = None
         self._address = None
 
-        if raw_script:
+        assert isinstance(mode, MiniWalletMode)
+        if mode == MiniWalletMode.RAW_OP_TRUE:
             self._scriptPubKey = bytes(CScript([OP_TRUE]))
-        elif use_p2pk:
+        elif mode == MiniWalletMode.RAW_P2PK:
             # use simple deterministic private key (k=1)
             self._priv_key = ECKey()
             self._priv_key.set((1).to_bytes(32, 'big'), True)
             pub_key = self._priv_key.get_pubkey()
             self._scriptPubKey = bytes(CScript([pub_key.get_bytes(), OP_CHECKSIG]))
-        else:
+        elif mode == MiniWalletMode.ADDRESS_OP_TRUE:
             self._address = ADDRESS_BCRT1_P2WSH_OP_TRUE
             self._scriptPubKey = hex_str_to_bytes(self._test_node.validateaddress(self._address)['scriptPubKey'])
 
