From 66c9936455fc0b79629686fe6b5737a020662e5d Mon Sep 17 00:00:00 2001
From: furszy <matiasfurszyfer@protonmail.com>
Date: Fri, 29 Sep 2023 17:23:48 -0300
Subject: [PATCH] bench: add coverage for wallet migration process

---
 src/bench/CMakeLists.txt       |  1 +
 src/bench/wallet_migration.cpp | 80 ++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+)
 create mode 100644 src/bench/wallet_migration.cpp

diff --git a/src/bench/CMakeLists.txt b/src/bench/CMakeLists.txt
index 8a52980e072..45b65100446 100644
--- a/src/bench/CMakeLists.txt
+++ b/src/bench/CMakeLists.txt
@@ -71,6 +71,7 @@ if(ENABLE_WALLET)
       wallet_create_tx.cpp
       wallet_loading.cpp
       wallet_ismine.cpp
+      wallet_migration.cpp
   )
   target_link_libraries(bench_bitcoin bitcoin_wallet)
 endif()
diff --git a/src/bench/wallet_migration.cpp b/src/bench/wallet_migration.cpp
new file mode 100644
index 00000000000..eff6c6b526d
--- /dev/null
+++ b/src/bench/wallet_migration.cpp
@@ -0,0 +1,80 @@
+// Copyright (c) 2024 The Bitcoin Core developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or https://www.opensource.org/licenses/mit-license.php.
+
+#include <bitcoin-build-config.h> // IWYU pragma: keep
+
+#include <bench/bench.h>
+#include <interfaces/chain.h>
+#include <node/context.h>
+#include <test/util/mining.h>
+#include <test/util/setup_common.h>
+#include <wallet/test/util.h>
+#include <wallet/context.h>
+#include <wallet/receive.h>
+#include <wallet/wallet.h>
+
+#include <optional>
+
+#if defined(USE_BDB) && defined(USE_SQLITE) // only enable benchmark when bdb and sqlite are enabled
+
+namespace wallet{
+
+static void WalletMigration(benchmark::Bench& bench)
+{
+    const auto test_setup = MakeNoLogFileContext<TestingSetup>();
+
+    WalletContext context;
+    context.args = &test_setup->m_args;
+    context.chain = test_setup->m_node.chain.get();
+
+    // Number of imported watch only addresses
+    int NUM_WATCH_ONLY_ADDR = 20;
+
+    // Setup legacy wallet
+    DatabaseOptions options;
+    options.use_unsafe_sync = true;
+    options.verify = false;
+    DatabaseStatus status;
+    bilingual_str error;
+    auto database = MakeWalletDatabase(fs::PathToString(test_setup->m_path_root / "legacy"), options, status, error);
+    uint64_t create_flags = 0;
+    auto wallet = TestLoadWallet(std::move(database), context, create_flags);
+
+    // Add watch-only addresses
+    std::vector<CScript> scripts_watch_only;
+    for (int w = 0; w < NUM_WATCH_ONLY_ADDR; ++w) {
+        CKey key = GenerateRandomKey();
+        LOCK(wallet->cs_wallet);
+        const CScript& script = scripts_watch_only.emplace_back(GetScriptForDestination(GetDestinationForKey(key.GetPubKey(), OutputType::LEGACY)));
+        bool res = wallet->ImportScriptPubKeys(strprintf("watch_%d", w), {script},
+                                    /*have_solving_data=*/false, /*apply_label=*/true, /*timestamp=*/1);
+        assert(res);
+    }
+
+    // Generate transactions and local addresses
+    for (int j = 0; j < 400; ++j) {
+        CMutableTransaction mtx;
+        mtx.vout.emplace_back(COIN, GetScriptForDestination(*Assert(wallet->GetNewDestination(OutputType::BECH32, strprintf("bench_%d", j)))));
+        mtx.vout.emplace_back(COIN, GetScriptForDestination(*Assert(wallet->GetNewDestination(OutputType::LEGACY, strprintf("legacy_%d", j)))));
+        mtx.vout.emplace_back(COIN, scripts_watch_only.at(j % NUM_WATCH_ONLY_ADDR));
+        mtx.vin.resize(2);
+        wallet->AddToWallet(MakeTransactionRef(mtx), TxStateInactive{}, /*update_wtx=*/nullptr, /*fFlushOnClose=*/false, /*rescanning_old_block=*/true);
+    }
+
+    // Unload so the migration process loads it
+    TestUnloadWallet(std::move(wallet));
+
+    bench.epochs(/*numEpochs=*/1).run([&] {
+        util::Result<MigrationResult> res = MigrateLegacyToDescriptor(fs::PathToString(test_setup->m_path_root / "legacy"), "", context);
+        assert(res);
+        assert(res->wallet);
+        assert(res->watchonly_wallet);
+    });
+}
+
+BENCHMARK(WalletMigration, benchmark::PriorityLevel::LOW);
+
+} // namespace wallet
+
+#endif // end USE_SQLITE && USE_BDB