From 51fc9241c08a00f1f407f1534853a5cddbbc0a23 Mon Sep 17 00:00:00 2001
From: James O'Beirne <james.obeirne@pm.me>
Date: Wed, 13 Apr 2022 15:16:26 -0400
Subject: [PATCH] test: allow on-disk coins and block tree dbs in tests

Used when testing cleanup of on-disk chainstate data for snapshot
testcases. Also necessary for simulating node restart in .cpp tests.
---
 src/test/util/chainstate.h     |  8 ++++++--
 src/test/util/setup_common.cpp | 22 ++++++++++++++++------
 src/test/util/setup_common.h   | 16 +++++++++++++---
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/src/test/util/chainstate.h b/src/test/util/chainstate.h
index 79bc8822158..0ca63810f33 100644
--- a/src/test/util/chainstate.h
+++ b/src/test/util/chainstate.h
@@ -30,7 +30,11 @@ const auto NoMalleation = [](AutoFile& file, node::SnapshotMetadata& meta){};
  */
 template<typename F = decltype(NoMalleation)>
 static bool
-CreateAndActivateUTXOSnapshot(TestingSetup* fixture, F malleation = NoMalleation, bool reset_chainstate = false)
+CreateAndActivateUTXOSnapshot(
+    TestingSetup* fixture,
+    F malleation = NoMalleation,
+    bool reset_chainstate = false,
+    bool in_memory_chainstate = false)
 {
     node::NodeContext& node = fixture->m_node;
     fs::path root = fixture->m_path_root;
@@ -88,7 +92,7 @@ CreateAndActivateUTXOSnapshot(TestingSetup* fixture, F malleation = NoMalleation
             0 == WITH_LOCK(node.chainman->GetMutex(), return node.chainman->ActiveHeight()));
     }
 
-    return node.chainman->ActivateSnapshot(auto_infile, metadata, /*in_memory=*/true);
+    return node.chainman->ActivateSnapshot(auto_infile, metadata, in_memory_chainstate);
 }
 
 
diff --git a/src/test/util/setup_common.cpp b/src/test/util/setup_common.cpp
index 74b055ee453..852c04e2d94 100644
--- a/src/test/util/setup_common.cpp
+++ b/src/test/util/setup_common.cpp
@@ -220,8 +220,14 @@ ChainTestingSetup::~ChainTestingSetup()
     m_node.chainman.reset();
 }
 
-TestingSetup::TestingSetup(const std::string& chainName, const std::vector<const char*>& extra_args)
-    : ChainTestingSetup(chainName, extra_args)
+TestingSetup::TestingSetup(
+    const std::string& chainName,
+    const std::vector<const char*>& extra_args,
+    const bool coins_db_in_memory,
+    const bool block_tree_db_in_memory)
+    : ChainTestingSetup(chainName, extra_args),
+      m_coins_db_in_memory(coins_db_in_memory),
+      m_block_tree_db_in_memory(block_tree_db_in_memory)
 {
     // Ideally we'd move all the RPC tests to the functional testing framework
     // instead of unit tests, but for now we need these here.
@@ -229,8 +235,8 @@ TestingSetup::TestingSetup(const std::string& chainName, const std::vector<const
 
     node::ChainstateLoadOptions options;
     options.mempool = Assert(m_node.mempool.get());
-    options.block_tree_db_in_memory = true;
-    options.coins_db_in_memory = true;
+    options.block_tree_db_in_memory = m_block_tree_db_in_memory;
+    options.coins_db_in_memory = m_coins_db_in_memory;
     options.reindex = node::fReindex;
     options.reindex_chainstate = m_args.GetBoolArg("-reindex-chainstate", false);
     options.prune = node::fPruneMode;
@@ -263,8 +269,12 @@ TestingSetup::TestingSetup(const std::string& chainName, const std::vector<const
     }
 }
 
-TestChain100Setup::TestChain100Setup(const std::string& chain_name, const std::vector<const char*>& extra_args)
-    : TestingSetup{chain_name, extra_args}
+TestChain100Setup::TestChain100Setup(
+        const std::string& chain_name,
+        const std::vector<const char*>& extra_args,
+        const bool coins_db_in_memory,
+        const bool block_tree_db_in_memory)
+    : TestingSetup{CBaseChainParams::REGTEST, extra_args, coins_db_in_memory, block_tree_db_in_memory}
 {
     SetMockTime(1598887952);
     constexpr std::array<unsigned char, 32> vchKey = {
diff --git a/src/test/util/setup_common.h b/src/test/util/setup_common.h
index 136ee1fd624..9fe2692c266 100644
--- a/src/test/util/setup_common.h
+++ b/src/test/util/setup_common.h
@@ -107,7 +107,14 @@ struct ChainTestingSetup : public BasicTestingSetup {
 /** Testing setup that configures a complete environment.
  */
 struct TestingSetup : public ChainTestingSetup {
-    explicit TestingSetup(const std::string& chainName = CBaseChainParams::MAIN, const std::vector<const char*>& extra_args = {});
+    bool m_coins_db_in_memory{true};
+    bool m_block_tree_db_in_memory{true};
+
+    explicit TestingSetup(
+        const std::string& chainName = CBaseChainParams::MAIN,
+        const std::vector<const char*>& extra_args = {},
+        const bool coins_db_in_memory = true,
+        const bool block_tree_db_in_memory = true);
 };
 
 /** Identical to TestingSetup, but chain set to regtest */
@@ -124,8 +131,11 @@ class CScript;
  * Testing fixture that pre-creates a 100-block REGTEST-mode block chain
  */
 struct TestChain100Setup : public TestingSetup {
-    TestChain100Setup(const std::string& chain_name = CBaseChainParams::REGTEST,
-                      const std::vector<const char*>& extra_args = {});
+    TestChain100Setup(
+        const std::string& chain_name = CBaseChainParams::REGTEST,
+        const std::vector<const char*>& extra_args = {},
+        const bool coins_db_in_memory = true,
+        const bool block_tree_db_in_memory = true);
 
     /**
      * Create a new block with just given transactions, coinbase paying to