4 files changed, 167 insertions, 51 deletions

diff --git a/build_msvc/libtest_util/libtest_util.vcxproj.in b/build_msvc/libtest_util/libtest_util.vcxproj.in
index b5e844010e..64cfa82dcc 100644
--- a/build_msvc/libtest_util/libtest_util.vcxproj.in
+++ b/build_msvc/libtest_util/libtest_util.vcxproj.in
@@ -8,6 +8,7 @@
     <ConfigurationType>StaticLibrary</ConfigurationType>
   </PropertyGroup>
   <ItemGroup>
+    <ClCompile Include="..\..\src\wallet\test\util.cpp" />
 @SOURCE_FILES@
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
diff --git a/src/wallet/test/util.cpp b/src/wallet/test/util.cpp
index b7bf312edf..7d0a814ed4 100644
--- a/src/wallet/test/util.cpp
+++ b/src/wallet/test/util.cpp
@@ -7,6 +7,7 @@
 #include <chain.h>
 #include <key.h>
 #include <key_io.h>
+#include <streams.h>
 #include <test/util/setup_common.h>
 #include <wallet/wallet.h>
 #include <wallet/walletdb.h>
@@ -79,4 +80,93 @@ CTxDestination getNewDestination(CWallet& w, OutputType output_type)
     return *Assert(w.GetNewDestination(output_type, ""));
 }
 
+DatabaseCursor::Status MockableCursor::Next(DataStream& key, DataStream& value)
+{
+    if (!m_pass) {
+        return Status::FAIL;
+    }
+    if (m_cursor == m_cursor_end) {
+        return Status::DONE;
+    }
+    const auto& [key_data, value_data] = *m_cursor;
+    key.write(key_data);
+    value.write(value_data);
+    m_cursor++;
+    return Status::MORE;
+}
+
+bool MockableBatch::ReadKey(DataStream&& key, DataStream& value)
+{
+    if (!m_pass) {
+        return false;
+    }
+    SerializeData key_data{key.begin(), key.end()};
+    const auto& it = m_records.find(key_data);
+    if (it == m_records.end()) {
+        return false;
+    }
+    value.write(it->second);
+    return true;
+}
+
+bool MockableBatch::WriteKey(DataStream&& key, DataStream&& value, bool overwrite)
+{
+    if (!m_pass) {
+        return false;
+    }
+    SerializeData key_data{key.begin(), key.end()};
+    SerializeData value_data{value.begin(), value.end()};
+    auto [it, inserted] = m_records.emplace(key_data, value_data);
+    if (!inserted && overwrite) { // Overwrite if requested
+        it->second = value_data;
+        inserted = true;
+    }
+    return inserted;
+}
+
+bool MockableBatch::EraseKey(DataStream&& key)
+{
+    if (!m_pass) {
+        return false;
+    }
+    SerializeData key_data{key.begin(), key.end()};
+    m_records.erase(key_data);
+    return true;
+}
+
+bool MockableBatch::HasKey(DataStream&& key)
+{
+    if (!m_pass) {
+        return false;
+    }
+    SerializeData key_data{key.begin(), key.end()};
+    return m_records.count(key_data) > 0;
+}
+
+bool MockableBatch::ErasePrefix(Span<const std::byte> prefix)
+{
+    if (!m_pass) {
+        return false;
+    }
+    auto it = m_records.begin();
+    while (it != m_records.end()) {
+        auto& key = it->first;
+        if (key.size() < prefix.size() || std::search(key.begin(), key.end(), prefix.begin(), prefix.end()) != key.begin()) {
+            it++;
+            continue;
+        }
+        it = m_records.erase(it);
+    }
+    return true;
+}
+
+std::unique_ptr<WalletDatabase> CreateMockableWalletDatabase(std::map<SerializeData, SerializeData> records)
+{
+    return std::make_unique<MockableDatabase>(records);
+}
+
+MockableDatabase& GetMockableDatabase(CWallet& wallet)
+{
+    return dynamic_cast<MockableDatabase&>(wallet.GetDatabase());
+}
 } // namespace wallet
diff --git a/src/wallet/test/util.h b/src/wallet/test/util.h
index d726517e21..92405107cf 100644
--- a/src/wallet/test/util.h
+++ b/src/wallet/test/util.h
@@ -6,6 +6,8 @@
 #define BITCOIN_WALLET_TEST_UTIL_H
 
 #include <script/standard.h>
+#include <wallet/db.h>
+
 #include <memory>
 
 class ArgsManager;
@@ -31,6 +33,78 @@ std::string getnewaddress(CWallet& w);
 /** Returns a new destination, of an specific type, from the wallet */
 CTxDestination getNewDestination(CWallet& w, OutputType output_type);
 
+class MockableCursor: public DatabaseCursor
+{
+public:
+    std::map<SerializeData, SerializeData>::const_iterator m_cursor;
+    std::map<SerializeData, SerializeData>::const_iterator m_cursor_end;
+    bool m_pass;
+
+    explicit MockableCursor(const std::map<SerializeData, SerializeData>& records, bool pass) : m_cursor(records.begin()), m_cursor_end(records.end()), m_pass(pass) {}
+    ~MockableCursor() {}
+
+    Status Next(DataStream& key, DataStream& value) override;
+};
+
+class MockableBatch : public DatabaseBatch
+{
+private:
+    std::map<SerializeData, SerializeData>& m_records;
+    bool m_pass;
+
+    bool ReadKey(DataStream&& key, DataStream& value) override;
+    bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite=true) override;
+    bool EraseKey(DataStream&& key) override;
+    bool HasKey(DataStream&& key) override;
+    bool ErasePrefix(Span<const std::byte> prefix) override;
+
+public:
+    explicit MockableBatch(std::map<SerializeData, SerializeData>& records, bool pass) : m_records(records), m_pass(pass) {}
+    ~MockableBatch() {}
+
+    void Flush() override {}
+    void Close() override {}
+
+    std::unique_ptr<DatabaseCursor> GetNewCursor() override
+    {
+        return std::make_unique<MockableCursor>(m_records, m_pass);
+    }
+    bool TxnBegin() override { return m_pass; }
+    bool TxnCommit() override { return m_pass; }
+    bool TxnAbort() override { return m_pass; }
+};
+
+/** A WalletDatabase whose contents and return values can be modified as needed for testing
+ **/
+class MockableDatabase : public WalletDatabase
+{
+public:
+    std::map<SerializeData, SerializeData> m_records;
+    bool m_pass{true};
+
+    MockableDatabase(std::map<SerializeData, SerializeData> records = {}) : WalletDatabase(), m_records(records) {}
+    ~MockableDatabase() {};
+
+    void Open() override {}
+    void AddRef() override {}
+    void RemoveRef() override {}
+
+    bool Rewrite(const char* pszSkip=nullptr) override { return m_pass; }
+    bool Backup(const std::string& strDest) const override { return m_pass; }
+    void Flush() override {}
+    void Close() override {}
+    bool PeriodicFlush() override { return m_pass; }
+    void IncrementUpdateCounter() override {}
+    void ReloadDbEnv() override {}
+
+    std::string Filename() override { return "mockable"; }
+    std::string Format() override { return "mock"; }
+    std::unique_ptr<DatabaseBatch> MakeBatch(bool flush_on_close = true) override { return std::make_unique<MockableBatch>(m_records, m_pass); }
+};
+
+std::unique_ptr<WalletDatabase> CreateMockableWalletDatabase(std::map<SerializeData, SerializeData> records = {});
+
+MockableDatabase& GetMockableDatabase(CWallet& wallet);
 } // namespace wallet
 
 #endif // BITCOIN_WALLET_TEST_UTIL_H
diff --git a/src/wallet/test/wallet_tests.cpp b/src/wallet/test/wallet_tests.cpp
index da4b553394..9cf045f065 100644
--- a/src/wallet/test/wallet_tests.cpp
+++ b/src/wallet/test/wallet_tests.cpp
@@ -951,62 +951,13 @@ BOOST_FIXTURE_TEST_CASE(ZapSelectTx, TestChain100Setup)
     TestUnloadWallet(std::move(wallet));
 }
 
-class FailCursor : public DatabaseCursor
-{
-public:
-    Status Next(DataStream& key, DataStream& value) override { return Status::FAIL; }
-};
-
-/** RAII class that provides access to a FailDatabase. Which fails if needed. */
-class FailBatch : public DatabaseBatch
-{
-private:
-    bool m_pass{true};
-    bool ReadKey(DataStream&& key, DataStream& value) override { return m_pass; }
-    bool WriteKey(DataStream&& key, DataStream&& value, bool overwrite = true) override { return m_pass; }
-    bool EraseKey(DataStream&& key) override { return m_pass; }
-    bool HasKey(DataStream&& key) override { return m_pass; }
-    bool ErasePrefix(Span<const std::byte> prefix) override { return m_pass; }
-
-public:
-    explicit FailBatch(bool pass) : m_pass(pass) {}
-    void Flush() override {}
-    void Close() override {}
-
-    std::unique_ptr<DatabaseCursor> GetNewCursor() override { return std::make_unique<FailCursor>(); }
-    bool TxnBegin() override { return false; }
-    bool TxnCommit() override { return false; }
-    bool TxnAbort() override { return false; }
-};
-
-/** A dummy WalletDatabase that does nothing, only fails if needed.**/
-class FailDatabase : public WalletDatabase
-{
-public:
-    bool m_pass{true}; // false when this db should fail
-
-    void Open() override {};
-    void AddRef() override {}
-    void RemoveRef() override {}
-    bool Rewrite(const char* pszSkip=nullptr) override { return true; }
-    bool Backup(const std::string& strDest) const override { return true; }
-    void Close() override {}
-    void Flush() override {}
-    bool PeriodicFlush() override { return true; }
-    void IncrementUpdateCounter() override { ++nUpdateCounter; }
-    void ReloadDbEnv() override {}
-    std::string Filename() override { return "faildb"; }
-    std::string Format() override { return "faildb"; }
-    std::unique_ptr<DatabaseBatch> MakeBatch(bool flush_on_close = true) override { return std::make_unique<FailBatch>(m_pass); }
-};
-
 /**
  * Checks a wallet invalid state where the inputs (prev-txs) of a new arriving transaction are not marked dirty,
  * while the transaction that spends them exist inside the in-memory wallet tx map (not stored on db due a db write failure).
  */
 BOOST_FIXTURE_TEST_CASE(wallet_sync_tx_invalid_state_test, TestingSetup)
 {
-    CWallet wallet(m_node.chain.get(), "", std::make_unique<FailDatabase>());
+    CWallet wallet(m_node.chain.get(), "", CreateMockableWalletDatabase());
     {
         LOCK(wallet.cs_wallet);
         wallet.SetWalletFlag(WALLET_FLAG_DESCRIPTORS);
@@ -1053,7 +1004,7 @@ BOOST_FIXTURE_TEST_CASE(wallet_sync_tx_invalid_state_test, TestingSetup)
     // 1) Make db always fail
     // 2) Try to add a transaction that spends the previously created transaction and
     //    verify that we are not moving forward if the wallet cannot store it
-    static_cast<FailDatabase&>(wallet.GetDatabase()).m_pass = false;
+    GetMockableDatabase(wallet).m_pass = false;
     mtx.vin.clear();
     mtx.vin.push_back(CTxIn(good_tx_id, 0));
     BOOST_CHECK_EXCEPTION(wallet.transactionAddedToMempool(MakeTransactionRef(mtx)),