Skip to content

Commit f81fad5

Browse files
committed
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 parent 61de64d commit f81fad5

File tree

3 files changed

+29
-29
lines changed

3 files changed

+29
-29
lines changed

test/functional/rpc_psbt.py

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616
CTxIn,
1717
CTxOut,
1818
MAX_BIP125_RBF_SEQUENCE,
19-
WITNESS_SCALE_FACTOR,
20-
ser_compact_size,
2119
)
2220
from test_framework.psbt import (
2321
PSBT,
@@ -42,6 +40,7 @@
4240
find_vout_for_address,
4341
)
4442
from test_framework.wallet_util import (
43+
calculate_input_weight,
4544
generate_keypair,
4645
get_generate_key,
4746
)
@@ -752,17 +751,9 @@ def test_psbt_input_keys(psbt_input, keys):
752751
input_idx = i
753752
break
754753
psbt_in = dec["inputs"][input_idx]
755-
# Calculate the input weight
756-
# (prevout + sequence + length of scriptSig + scriptsig) * WITNESS_SCALE_FACTOR + len of num scriptWitness stack items + (length of stack item + stack item) * N stack items
757-
# Note that occasionally this weight estimate may be slightly larger or smaller than the real weight
758-
# as sometimes ECDSA signatures are one byte shorter than expected with a probability of 1/128
759-
len_scriptsig = len(psbt_in["final_scriptSig"]["hex"]) // 2 if "final_scriptSig" in psbt_in else 0
760-
len_scriptsig += len(ser_compact_size(len_scriptsig))
761-
len_scriptwitness = (sum([(len(x) // 2) + len(ser_compact_size(len(x) // 2)) for x in psbt_in["final_scriptwitness"]]) + len(ser_compact_size(len(psbt_in["final_scriptwitness"])))) if "final_scriptwitness" in psbt_in else 0
762-
len_prevout_txid = 32
763-
len_prevout_index = 4
764-
len_sequence = 4
765-
input_weight = ((len_prevout_txid + len_prevout_index + len_sequence + len_scriptsig) * WITNESS_SCALE_FACTOR) + len_scriptwitness
754+
scriptsig_hex = psbt_in["final_scriptSig"]["hex"] if "final_scriptSig" in psbt_in else ""
755+
witness_stack_hex = psbt_in["final_scriptwitness"] if "final_scriptwitness" in psbt_in else None
756+
input_weight = calculate_input_weight(scriptsig_hex, witness_stack_hex)
766757
low_input_weight = input_weight // 2
767758
high_input_weight = input_weight * 2
768759

test/functional/test_framework/wallet_util.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@
1515
script_to_p2wsh,
1616
)
1717
from test_framework.key import ECKey
18+
from test_framework.messages import (
19+
CTxIn,
20+
CTxInWitness,
21+
WITNESS_SCALE_FACTOR,
22+
)
1823
from test_framework.script_util import (
1924
key_to_p2pkh_script,
2025
key_to_p2wpkh_script,
@@ -123,6 +128,19 @@ def generate_keypair(compressed=True, wif=False):
123128
privkey = bytes_to_wif(privkey.get_bytes(), compressed)
124129
return privkey, pubkey
125130

131+
def calculate_input_weight(scriptsig_hex, witness_stack_hex=None):
132+
"""Given a scriptSig and a list of witness stack items for an input in hex format,
133+
calculate the total input weight. If the input has no witness data,
134+
`witness_stack_hex` can be set to None."""
135+
tx_in = CTxIn(scriptSig=bytes.fromhex(scriptsig_hex))
136+
witness_size = 0
137+
if witness_stack_hex is not None:
138+
tx_inwit = CTxInWitness()
139+
for witness_item_hex in witness_stack_hex:
140+
tx_inwit.scriptWitness.stack.append(bytes.fromhex(witness_item_hex))
141+
witness_size = len(tx_inwit.serialize())
142+
return len(tx_in.serialize()) * WITNESS_SCALE_FACTOR + witness_size
143+
126144
class WalletUnlock():
127145
"""
128146
A context manager for unlocking a wallet with a passphrase and automatically locking it afterward.

test/functional/wallet_send.py

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,6 @@
99

1010
from test_framework.authproxy import JSONRPCException
1111
from test_framework.descriptors import descsum_create
12-
from test_framework.messages import (
13-
ser_compact_size,
14-
WITNESS_SCALE_FACTOR,
15-
)
1612
from test_framework.test_framework import BitcoinTestFramework
1713
from test_framework.util import (
1814
assert_equal,
@@ -21,7 +17,10 @@
2117
assert_raises_rpc_error,
2218
count_bytes,
2319
)
24-
from test_framework.wallet_util import generate_keypair
20+
from test_framework.wallet_util import (
21+
calculate_input_weight,
22+
generate_keypair,
23+
)
2524

2625

2726
class WalletSendTest(BitcoinTestFramework):
@@ -543,17 +542,9 @@ def run_test(self):
543542
input_idx = i
544543
break
545544
psbt_in = dec["inputs"][input_idx]
546-
# Calculate the input weight
547-
# (prevout + sequence + length of scriptSig + scriptsig) * WITNESS_SCALE_FACTOR + len of num scriptWitness stack items + (length of stack item + stack item) * N stack items
548-
# Note that occasionally this weight estimate may be slightly larger or smaller than the real weight
549-
# as sometimes ECDSA signatures are one byte shorter than expected with a probability of 1/128
550-
len_scriptsig = len(psbt_in["final_scriptSig"]["hex"]) // 2 if "final_scriptSig" in psbt_in else 0
551-
len_scriptsig += len(ser_compact_size(len_scriptsig))
552-
len_scriptwitness = (sum([(len(x) // 2) + len(ser_compact_size(len(x) // 2)) for x in psbt_in["final_scriptwitness"]]) + len(ser_compact_size(len(psbt_in["final_scriptwitness"])))) if "final_scriptwitness" in psbt_in else 0
553-
len_prevout_txid = 32
554-
len_prevout_index = 4
555-
len_sequence = 4
556-
input_weight = ((len_prevout_txid + len_prevout_index + len_sequence + len_scriptsig) * WITNESS_SCALE_FACTOR) + len_scriptwitness
545+
scriptsig_hex = psbt_in["final_scriptSig"]["hex"] if "final_scriptSig" in psbt_in else ""
546+
witness_stack_hex = psbt_in["final_scriptwitness"] if "final_scriptwitness" in psbt_in else None
547+
input_weight = calculate_input_weight(scriptsig_hex, witness_stack_hex)
557548

558549
# Input weight error conditions
559550
assert_raises_rpc_error(

0 commit comments

Comments
 (0)