Add ChainNotifier and define ChainListener trait

[rust-lightning] / lightning-block-sync / src / test_utils.rs

diff --git a/lightning-block-sync/src/test_utils.rs b/lightning-block-sync/src/test_utils.rs
index efd34f74ee5b3cb219566f1dcad5a8d4d9d57315..70a8982a0efb746a303bdcd530ec536c322e0be3 100644 (file)
--- a/lightning-block-sync/src/test_utils.rs
+++ b/lightning-block-sync/src/test_utils.rs
@@ -1,4 +1,4 @@
-use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSource, BlockSourceError};
+use crate::{AsyncBlockSourceResult, BlockHeaderData, BlockSource, BlockSourceError, ChainListener, UnboundedCache};
 use crate::poll::{Validate, ValidatedBlockHeader};
 
 use bitcoin::blockdata::block::{Block, BlockHeader};
@@ -7,9 +7,12 @@ use bitcoin::hash_types::BlockHash;
 use bitcoin::network::constants::Network;
 use bitcoin::util::uint::Uint256;
 
+use std::collections::VecDeque;
+
 #[derive(Default)]
 pub struct Blockchain {
        pub blocks: Vec<Block>,
+       without_blocks: Option<std::ops::RangeFrom<usize>>,
        without_headers: bool,
        malformed_headers: bool,
 }
@@ -46,6 +49,10 @@ impl Blockchain {
                self
        }
 
+       pub fn without_blocks(self, range: std::ops::RangeFrom<usize>) -> Self {
+               Self { without_blocks: Some(range), ..self }
+       }
+
        pub fn without_headers(self) -> Self {
                Self { without_headers: true, ..self }
        }
@@ -54,6 +61,18 @@ impl Blockchain {
                Self { malformed_headers: true, ..self }
        }
 
+       pub fn fork_at_height(&self, height: usize) -> Self {
+               assert!(height + 1 < self.blocks.len());
+               let mut blocks = self.blocks.clone();
+               let mut prev_blockhash = blocks[height].block_hash();
+               for block in blocks.iter_mut().skip(height + 1) {
+                       block.header.prev_blockhash = prev_blockhash;
+                       block.header.nonce += 1;
+                       prev_blockhash = block.block_hash();
+               }
+               Self { blocks, without_blocks: None, ..*self }
+       }
+
        pub fn at_height(&self, height: usize) -> ValidatedBlockHeader {
                let block_header = self.at_height_unvalidated(height);
                let block_hash = self.blocks[height].block_hash();
@@ -78,6 +97,16 @@ impl Blockchain {
        pub fn disconnect_tip(&mut self) -> Option<Block> {
                self.blocks.pop()
        }
+
+       pub fn header_cache(&self, heights: std::ops::RangeInclusive<usize>) -> UnboundedCache {
+               let mut cache = UnboundedCache::new();
+               for i in heights {
+                       let value = self.at_height(i);
+                       let key = value.header.block_hash();
+                       assert!(cache.insert(key, value).is_none());
+               }
+               cache
+       }
 }
 
 impl BlockSource for Blockchain {
@@ -103,8 +132,14 @@ impl BlockSource for Blockchain {
 
        fn get_block<'a>(&'a mut self, header_hash: &'a BlockHash) -> AsyncBlockSourceResult<'a, Block> {
                Box::pin(async move {
-                       for block in self.blocks.iter() {
+                       for (height, block) in self.blocks.iter().enumerate() {
                                if block.header.block_hash() == *header_hash {
+                                       if let Some(without_blocks) = &self.without_blocks {
+                                               if without_blocks.contains(&height) {
+                                                       return Err(BlockSourceError::persistent("block not found"));
+                                               }
+                                       }
+
                                        return Ok(block.clone());
                                }
                        }
@@ -124,3 +159,67 @@ impl BlockSource for Blockchain {
                })
        }
 }
+
+pub struct MockChainListener {
+       expected_blocks_connected: VecDeque<BlockHeaderData>,
+       expected_blocks_disconnected: VecDeque<BlockHeaderData>,
+}
+
+impl MockChainListener {
+       pub fn new() -> Self {
+               Self {
+                       expected_blocks_connected: VecDeque::new(),
+                       expected_blocks_disconnected: VecDeque::new(),
+               }
+       }
+
+       pub fn expect_block_connected(mut self, block: BlockHeaderData) -> Self {
+               self.expected_blocks_connected.push_back(block);
+               self
+       }
+
+       pub fn expect_block_disconnected(mut self, block: BlockHeaderData) -> Self {
+               self.expected_blocks_disconnected.push_back(block);
+               self
+       }
+}
+
+impl ChainListener for MockChainListener {
+       fn block_connected(&mut self, block: &Block, height: u32) {
+               match self.expected_blocks_connected.pop_front() {
+                       None => {
+                               panic!("Unexpected block connected: {:?}", block.block_hash());
+                       },
+                       Some(expected_block) => {
+                               assert_eq!(block.block_hash(), expected_block.header.block_hash());
+                               assert_eq!(height, expected_block.height);
+                       },
+               }
+       }
+
+       fn block_disconnected(&mut self, header: &BlockHeader, height: u32) {
+               match self.expected_blocks_disconnected.pop_front() {
+                       None => {
+                               panic!("Unexpected block disconnected: {:?}", header.block_hash());
+                       },
+                       Some(expected_block) => {
+                               assert_eq!(header.block_hash(), expected_block.header.block_hash());
+                               assert_eq!(height, expected_block.height);
+                       },
+               }
+       }
+}
+
+impl Drop for MockChainListener {
+       fn drop(&mut self) {
+               if std::thread::panicking() {
+                       return;
+               }
+               if !self.expected_blocks_connected.is_empty() {
+                       panic!("Expected blocks connected: {:?}", self.expected_blocks_connected);
+               }
+               if !self.expected_blocks_disconnected.is_empty() {
+                       panic!("Expected blocks disconnected: {:?}", self.expected_blocks_disconnected);
+               }
+       }
+}