2 files changed, 62 insertions, 6 deletions

diff --git a/src/wallet/test/spend_tests.cpp b/src/wallet/test/spend_tests.cpp
index b2a0697c21..334bd5b8bc 100644
--- a/src/wallet/test/spend_tests.cpp
+++ b/src/wallet/test/spend_tests.cpp
@@ -63,5 +63,56 @@ BOOST_FIXTURE_TEST_CASE(SubtractFee, TestChain100Setup)
     BOOST_CHECK_EQUAL(fee, check_tx(fee + 123));
 }
 
+static void TestFillInputToWeight(int64_t additional_weight, std::vector<int64_t> expected_stack_sizes)
+{
+    static const int64_t EMPTY_INPUT_WEIGHT = GetTransactionInputWeight(CTxIn());
+
+    CTxIn input;
+    int64_t target_weight = EMPTY_INPUT_WEIGHT + additional_weight;
+    BOOST_CHECK(FillInputToWeight(input, target_weight));
+    BOOST_CHECK_EQUAL(GetTransactionInputWeight(input), target_weight);
+    BOOST_CHECK_EQUAL(input.scriptWitness.stack.size(), expected_stack_sizes.size());
+    for (unsigned int i = 0; i < expected_stack_sizes.size(); ++i) {
+        BOOST_CHECK_EQUAL(input.scriptWitness.stack[i].size(), expected_stack_sizes[i]);
+    }
+}
+
+BOOST_FIXTURE_TEST_CASE(FillInputToWeightTest, BasicTestingSetup)
+{
+    {
+        // Less than or equal minimum of 165 should not add any witness data
+        CTxIn input;
+        BOOST_CHECK(!FillInputToWeight(input, -1));
+        BOOST_CHECK_EQUAL(GetTransactionInputWeight(input), 165);
+        BOOST_CHECK_EQUAL(input.scriptWitness.stack.size(), 0);
+        BOOST_CHECK(!FillInputToWeight(input, 0));
+        BOOST_CHECK_EQUAL(GetTransactionInputWeight(input), 165);
+        BOOST_CHECK_EQUAL(input.scriptWitness.stack.size(), 0);
+        BOOST_CHECK(!FillInputToWeight(input, 164));
+        BOOST_CHECK_EQUAL(GetTransactionInputWeight(input), 165);
+        BOOST_CHECK_EQUAL(input.scriptWitness.stack.size(), 0);
+        BOOST_CHECK(FillInputToWeight(input, 165));
+        BOOST_CHECK_EQUAL(GetTransactionInputWeight(input), 165);
+        BOOST_CHECK_EQUAL(input.scriptWitness.stack.size(), 0);
+    }
+
+    // Make sure we can add at least one weight
+    TestFillInputToWeight(1, {0});
+
+    // 1 byte compact size uint boundary
+    TestFillInputToWeight(252, {251});
+    TestFillInputToWeight(253, {83, 168});
+    TestFillInputToWeight(262, {86, 174});
+    TestFillInputToWeight(263, {260});
+
+    // 3 byte compact size uint boundary
+    TestFillInputToWeight(65535, {65532});
+    TestFillInputToWeight(65536, {21842, 43688});
+    TestFillInputToWeight(65545, {21845, 43694});
+    TestFillInputToWeight(65546, {65541});
+
+    // Note: We don't test the next boundary because of memory allocation constraints.
+}
+
 BOOST_AUTO_TEST_SUITE_END()
 } // namespace wallet
diff --git a/src/wallet/test/wallet_tests.cpp b/src/wallet/test/wallet_tests.cpp
index 55428f43a7..8ef0d46c4f 100644
--- a/src/wallet/test/wallet_tests.cpp
+++ b/src/wallet/test/wallet_tests.cpp
@@ -140,11 +140,13 @@ BOOST_FIXTURE_TEST_CASE(scan_for_wallet_transactions, TestChain100Setup)
     }
 
     // Prune the older block file.
+    int file_number;
     {
         LOCK(cs_main);
-        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(oldTip->GetBlockPos().nFile);
+        file_number = oldTip->GetBlockPos().nFile;
+        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(file_number);
     }
-    UnlinkPrunedFiles({oldTip->GetBlockPos().nFile});
+    UnlinkPrunedFiles({file_number});
 
     // Verify ScanForWalletTransactions only picks transactions in the new block
     // file.
@@ -169,9 +171,10 @@ BOOST_FIXTURE_TEST_CASE(scan_for_wallet_transactions, TestChain100Setup)
     // Prune the remaining block file.
     {
         LOCK(cs_main);
-        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(newTip->GetBlockPos().nFile);
+        file_number = newTip->GetBlockPos().nFile;
+        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(file_number);
     }
-    UnlinkPrunedFiles({newTip->GetBlockPos().nFile});
+    UnlinkPrunedFiles({file_number});
 
     // Verify ScanForWalletTransactions scans no blocks.
     {
@@ -202,11 +205,13 @@ BOOST_FIXTURE_TEST_CASE(importmulti_rescan, TestChain100Setup)
     CBlockIndex* newTip = m_node.chainman->ActiveChain().Tip();
 
     // Prune the older block file.
+    int file_number;
     {
         LOCK(cs_main);
-        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(oldTip->GetBlockPos().nFile);
+        file_number = oldTip->GetBlockPos().nFile;
+        Assert(m_node.chainman)->m_blockman.PruneOneBlockFile(file_number);
     }
-    UnlinkPrunedFiles({oldTip->GetBlockPos().nFile});
+    UnlinkPrunedFiles({file_number});
 
     // Verify importmulti RPC returns failure for a key whose creation time is
     // before the missing block, and success for a key whose creation time is