From 33c6245ac1ecdfe25b1ee4fd9e93c43393634ae3 Mon Sep 17 00:00:00 2001
From: Andrew Chow <github@achow101.com>
Date: Mon, 12 Dec 2022 14:54:01 -0500
Subject: Introduce MockableDatabase for wallet unit tests

MockableDatabase is a WalletDatabase that allows us to interact with the
records to change them independently from the wallet, as well as
changing the return values from within the tests. This will give us
greater flexibility in testing the wallet.
---
 src/wallet/test/util.cpp         | 90 ++++++++++++++++++++++++++++++++++++++++
 src/wallet/test/util.h           | 74 +++++++++++++++++++++++++++++++++
 src/wallet/test/wallet_tests.cpp | 53 +----------------------
 3 files changed, 166 insertions(+), 51 deletions(-)

(limited to 'src/wallet')

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)),
-- 
cgit v1.2.3