From 6d91cb781c30966963f28e7577c7aa3829fa9390 Mon Sep 17 00:00:00 2001
From: Sebastian Falbesoner <sebastian.falbesoner@gmail.com>
Date: Mon, 1 Apr 2024 15:01:38 +0200
Subject: [PATCH] test: add unit tests for `calculate_input_weight`

---
 test/functional/test_framework/wallet_util.py | 40 +++++++++++++++++++
 test/functional/test_runner.py                |  1 +
 2 files changed, 41 insertions(+)

diff --git a/test/functional/test_framework/wallet_util.py b/test/functional/test_framework/wallet_util.py
index 8104701ff35..d30b00f4a76 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,
@@ -159,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)
diff --git a/test/functional/test_runner.py b/test/functional/test_runner.py
index 3f6e47d4103..32b55813a8c 100755
--- a/test/functional/test_runner.py
+++ b/test/functional/test_runner.py
@@ -85,6 +85,7 @@ TEST_FRAMEWORK_MODULES = [
     "crypto.ripemd160",
     "script",
     "segwit_addr",
+    "wallet_util",
 ]
 
 EXTENDED_SCRIPTS = [