Update wallet last processed index in every unit test that needs it

furszy · furszy · commit 6320efbbef23 · 2021-03-10T16:24:14.000-03:00
diff --git a/src/wallet/test/wallet_shielded_balances_tests.cpp b/src/wallet/test/wallet_shielded_balances_tests.cpp
@@ -291,7 +291,7 @@ struct FakeBlock
     CBlockIndex* pindex;
 };
 
-FakeBlock SimpleFakeMine(CWalletTx& wtx, SaplingMerkleTree& currentTree)
+FakeBlock SimpleFakeMine(CWalletTx& wtx, SaplingMerkleTree& currentTree, CWallet& wallet)
 {
     FakeBlock fakeBlock;
     fakeBlock.block.nVersion = 8;
@@ -306,6 +306,7 @@ FakeBlock SimpleFakeMine(CWalletTx& wtx, SaplingMerkleTree& currentTree)
     fakeBlock.pindex->phashBlock = &mapBlockIndex.find(fakeBlock.block.GetHash())->first;
     chainActive.SetTip(fakeBlock.pindex);
     BOOST_CHECK(chainActive.Contains(fakeBlock.pindex));
+    WITH_LOCK(wallet.cs_wallet, wallet.SetLastBlockProcessed(fakeBlock.pindex));
     wtx.SetConf(CWalletTx::Status::CONFIRMED, fakeBlock.pindex->GetBlockHash(), 0);
     return fakeBlock;
 }
@@ -347,7 +348,7 @@ BOOST_AUTO_TEST_CASE(GetShieldedAvailableCredit)
 
     // 2) Confirm the tx
     SaplingMerkleTree tree;
-    FakeBlock fakeBlock = SimpleFakeMine(wtxUpdated, tree);
+    FakeBlock fakeBlock = SimpleFakeMine(wtxUpdated, tree, wallet);
     // Simulate receiving a new block and updating the witnesses/nullifiers
     wallet.IncrementNoteWitnesses(fakeBlock.pindex, &fakeBlock.block, tree);
     wallet.GetSaplingScriptPubKeyMan()->UpdateSaplingNullifierNoteMapForBlock(&fakeBlock.block);
diff --git a/src/wallet/test/wallet_tests.cpp b/src/wallet/test/wallet_tests.cpp
@@ -320,7 +320,7 @@ void removeTxFromMempool(CWalletTx& wtx)
 /**
  * Mimic block creation.
  */
-CBlockIndex* SimpleFakeMine(CWalletTx& wtx, CBlockIndex* pprev = nullptr)
+CBlockIndex* SimpleFakeMine(CWalletTx& wtx, CWallet &wallet, CBlockIndex* pprev = nullptr)
 {
     CBlock block;
     block.vtx.emplace_back(wtx.tx);
@@ -332,6 +332,7 @@ CBlockIndex* SimpleFakeMine(CWalletTx& wtx, CBlockIndex* pprev = nullptr)
     fakeIndex->phashBlock = &mapBlockIndex.find(block.GetHash())->first;
     chainActive.SetTip(fakeIndex);
     BOOST_CHECK(chainActive.Contains(fakeIndex));
+    WITH_LOCK(wallet.cs_wallet, wallet.SetLastBlockProcessed(fakeIndex));
     wtx.SetConf(CWalletTx::Status::CONFIRMED, fakeIndex->GetBlockHash(), 0);
     removeTxFromMempool(wtx);
     wtx.fInMempool = false;
@@ -442,7 +443,7 @@ BOOST_AUTO_TEST_CASE(cached_balances_tests)
     BOOST_CHECK_EQUAL(wallet.GetUnconfirmedBalance(), nCredit);
 
     // 2) Confirm tx and verify
-    SimpleFakeMine(wtxCredit);
+    SimpleFakeMine(wtxCredit, wallet);
     BOOST_CHECK_EQUAL(wallet.GetUnconfirmedBalance(), 0);
     BOOST_CHECK_EQUAL(wtxCredit.GetAvailableCredit(), nCredit);
 
diff --git a/src/wallet/wallet.h b/src/wallet/wallet.h
@@ -644,8 +644,22 @@ class CWallet : public CCryptoKeyStore, public CValidationInterface
     bool IsHDEnabled() const;
     //! Whether the wallet supports Sapling or not //
     bool IsSaplingUpgradeEnabled() const;
-    //! Return the height of the last processed block
-    int GetLastBlockHeight() const { return m_last_block_processed_height; }
+
+    /** Get last block processed height */
+    int GetLastBlockHeight() const EXCLUSIVE_LOCKS_REQUIRED(cs_wallet)
+    {
+        AssertLockHeld(cs_wallet);
+        assert(m_last_block_processed_height >= 0);
+        return m_last_block_processed_height;
+    };
+    /** Set last block processed height, currently only use in unit test */
+    void SetLastBlockProcessed(const CBlockIndex* pindex) EXCLUSIVE_LOCKS_REQUIRED(cs_wallet)
+    {
+        AssertLockHeld(cs_wallet);
+        m_last_block_processed_height = pindex->nHeight;
+        m_last_block_processed = pindex->GetBlockHash();
+        m_last_block_processed_time = pindex->GetBlockTime();
+    };
 
     /* SPKM Helpers */
     const CKeyingMaterial& GetEncryptionKey() const;
