test: introduce and use `calculate_input_weight` helper

Rather than manually estimating an input's weight by adding up all the
involved components (fixed-size skeleton, compact-serialized lengths,
and the actual scriptSig / witness stack items) we can simply take use
of the serialization classes `CTxIn` / `CTxInWitness` instead, to
achieve the same with significantly less code.

The new helper is used in the functional tests rpc_psbt.py and
wallet_send.py, where the previous manual estimation code was
duplicated.

1 files changed, 18 insertions, 0 deletions

diff --git a/test/functional/test_framework/wallet_util.py b/test/functional/test_framework/wallet_util.py
index 44811918bf..8104701ff3 100755
--- a/test/functional/test_framework/wallet_util.py
+++ b/test/functional/test_framework/wallet_util.py
@@ -15,6 +15,11 @@ from test_framework.address import (
     script_to_p2wsh,
 )
 from test_framework.key import ECKey
+from test_framework.messages import (
+    CTxIn,
+    CTxInWitness,
+    WITNESS_SCALE_FACTOR,
+)
 from test_framework.script_util import (
     key_to_p2pkh_script,
     key_to_p2wpkh_script,
@@ -123,6 +128,19 @@ def generate_keypair(compressed=True, wif=False):
         privkey = bytes_to_wif(privkey.get_bytes(), compressed)
     return privkey, pubkey
 
+def calculate_input_weight(scriptsig_hex, witness_stack_hex=None):
+    """Given a scriptSig and a list of witness stack items for an input in hex format,
+       calculate the total input weight. If the input has no witness data,
+       `witness_stack_hex` can be set to None."""
+    tx_in = CTxIn(scriptSig=bytes.fromhex(scriptsig_hex))
+    witness_size = 0
+    if witness_stack_hex is not None:
+        tx_inwit = CTxInWitness()
+        for witness_item_hex in witness_stack_hex:
+            tx_inwit.scriptWitness.stack.append(bytes.fromhex(witness_item_hex))
+        witness_size = len(tx_inwit.serialize())
+    return len(tx_in.serialize()) * WITNESS_SCALE_FACTOR + witness_size
+
 class WalletUnlock():
     """
     A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.