[block encodings] Make CheckBlock mockable for PartiallyDownloadedBlock

2 files changed, 12 insertions, 1 deletions

diff --git a/src/blockencodings.cpp b/src/blockencodings.cpp
index bcb86d75cc..c83e893c0e 100644
--- a/src/blockencodings.cpp
+++ b/src/blockencodings.cpp
@@ -197,7 +197,8 @@ ReadStatus PartiallyDownloadedBlock::FillBlock(CBlock& block, const std::vector<
         return READ_STATUS_INVALID;
 
     BlockValidationState state;
-    if (!CheckBlock(block, state, Params().GetConsensus())) {
+    CheckBlockFn check_block = m_check_block_mock ? m_check_block_mock : CheckBlock;
+    if (!check_block(block, state, Params().GetConsensus(), /*fCheckPoW=*/true, /*fCheckMerkleRoot=*/true)) {
         // TODO: We really want to just check merkle tree manually here,
         // but that is expensive, and CheckBlock caches a block's
         // "checked-status" (in the CBlock?). CBlock should be able to
diff --git a/src/blockencodings.h b/src/blockencodings.h
index e60c1e3db4..7207ff1ae2 100644
--- a/src/blockencodings.h
+++ b/src/blockencodings.h
@@ -7,8 +7,13 @@
 
 #include <primitives/block.h>
 
+#include <functional>
 
 class CTxMemPool;
+class BlockValidationState;
+namespace Consensus {
+struct Params;
+};
 
 // Transaction compression schemes for compact block relay can be introduced by writing
 // an actual formatter here.
@@ -129,6 +134,11 @@ protected:
     const CTxMemPool* pool;
 public:
     CBlockHeader header;
+
+    // Can be overriden for testing
+    using CheckBlockFn = std::function<bool(const CBlock&, BlockValidationState&, const Consensus::Params&, bool, bool)>;
+    CheckBlockFn m_check_block_mock{nullptr};
+
     explicit PartiallyDownloadedBlock(CTxMemPool* poolIn) : pool(poolIn) {}
 
     // extra_txn is a list of extra transactions to look at, in <witness hash, reference> form