Merge bitcoin/bitcoin#29777: test: refactor: introduce and use `calculate_input_weight` helper

6d91cb781c30966963f28e7577c7aa3829fa9390 test: add unit tests for `calculate_input_weight` (Sebastian Falbesoner)
f81fad5e0f3be1f7aed59f9da00396c75c2a6406 test: introduce and use `calculate_input_weight` helper (Sebastian Falbesoner)

Pull request description:

  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. Unit tests are added in the second commit.

ACKs for top commit:
  kevkevinpal:
    tACK [6d91cb7](https://github.com/bitcoin/bitcoin/pull/29777/commits/6d91cb781c30966963f28e7577c7aa3829fa9390)
  QureshiFaisal:
    tACK [6d91cb7](https://github.com/bitcoin/bitcoin/pull/29777/commits/6d91cb781c30966963f28e7577c7aa3829fa9390)
  achow101:
    ACK 6d91cb781c30966963f28e7577c7aa3829fa9390
  AngusP:
    tACK 6d91cb781c30966963f28e7577c7aa3829fa9390
  rkrux:
    tACK [6d91cb7](https://github.com/bitcoin/bitcoin/pull/29777/commits/6d91cb781c30966963f28e7577c7aa3829fa9390)

Tree-SHA512: 04424e4d94d0e13745a9c11df2dd3697c98552bbb0e792c4af67ecbb66060adc3cc0cefc202cdee2d9db0baf85b8bedf2eb339ac4b316d986b5f10f6b70c5a33

1 files changed, 58 insertions, 0 deletions

diff --git a/test/functional/test_framework/wallet_util.py b/test/functional/test_framework/wallet_util.py
index 44811918bf..d30b00f4a7 100755
--- a/test/functional/test_framework/wallet_util.py
+++ b/test/functional/test_framework/wallet_util.py
@@ -4,6 +4,7 @@
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 """Useful util functions for testing the wallet"""
 from collections import namedtuple
+import unittest
 
 from test_framework.address import (
     byte_to_base58,
@@ -15,6 +16,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 +129,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.
@@ -141,3 +160,42 @@ class WalletUnlock():
     def __exit__(self, *args):
         _ = args
         self.wallet.walletlock()
+
+
+class TestFrameworkWalletUtil(unittest.TestCase):
+    def test_calculate_input_weight(self):
+        SKELETON_BYTES = 32 + 4 + 4  # prevout-txid, prevout-index, sequence
+        SMALL_LEN_BYTES = 1  # bytes needed for encoding scriptSig / witness item lenghts < 253
+        LARGE_LEN_BYTES = 3  # bytes needed for encoding scriptSig / witness item lengths >= 253
+
+        # empty scriptSig, no witness
+        self.assertEqual(calculate_input_weight(""),
+                         (SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR)
+        self.assertEqual(calculate_input_weight("", None),
+                         (SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR)
+        # small scriptSig, no witness
+        scriptSig_small = "00"*252
+        self.assertEqual(calculate_input_weight(scriptSig_small, None),
+                         (SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR)
+        # small scriptSig, empty witness stack
+        self.assertEqual(calculate_input_weight(scriptSig_small, []),
+                         (SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR + SMALL_LEN_BYTES)
+        # large scriptSig, no witness
+        scriptSig_large = "00"*253
+        self.assertEqual(calculate_input_weight(scriptSig_large, None),
+                         (SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR)
+        # large scriptSig, empty witness stack
+        self.assertEqual(calculate_input_weight(scriptSig_large, []),
+                         (SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR + SMALL_LEN_BYTES)
+        # empty scriptSig, 5 small witness stack items
+        self.assertEqual(calculate_input_weight("", ["00", "11", "22", "33", "44"]),
+                         ((SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 5 * SMALL_LEN_BYTES + 5)
+        # empty scriptSig, 253 small witness stack items
+        self.assertEqual(calculate_input_weight("", ["00"]*253),
+                         ((SKELETON_BYTES + SMALL_LEN_BYTES) * WITNESS_SCALE_FACTOR) + LARGE_LEN_BYTES + 253 * SMALL_LEN_BYTES + 253)
+        # small scriptSig, 3 large witness stack items
+        self.assertEqual(calculate_input_weight(scriptSig_small, ["00"*253]*3),
+                         ((SKELETON_BYTES + SMALL_LEN_BYTES + 252) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 3 * LARGE_LEN_BYTES + 3*253)
+        # large scriptSig, 3 large witness stack items
+        self.assertEqual(calculate_input_weight(scriptSig_large, ["00"*253]*3),
+                         ((SKELETON_BYTES + LARGE_LEN_BYTES + 253) * WITNESS_SCALE_FACTOR) + SMALL_LEN_BYTES + 3 * LARGE_LEN_BYTES + 3*253)