Restore and update mempool tests

zebrad/src/components/mempool/storage/tests/vectors.rs

@@ -18,6 +18,9 @@ use crate::components::mempool::{
 /// so we use a large enough value that will never be reached in the tests.
 const EVICTION_MEMORY_TIME: Duration = Duration::from_secs(60 * 60);
 
+/// Transaction count used in some tests to derive the mempool test size.
+const MEMPOOL_TX_COUNT: usize = 4;
+
 #[test]
 fn mempool_storage_crud_exact_mainnet() {
     zebra_test::init();
@@ -50,6 +53,101 @@ fn mempool_storage_crud_exact_mainnet() {
     assert!(!storage.contains_transaction_exact(&unmined_tx.transaction.id));
 }
 
+#[test]
+fn mempool_storage_basic() -> Result<()> {
+    zebra_test::init();
+
+    // Test multiple times to catch intermittent bugs since eviction is randomized
+    for _ in 0..10 {
+        mempool_storage_basic_for_network(Network::Mainnet)?;
+        mempool_storage_basic_for_network(Network::Testnet)?;
+    }
+
+    Ok(())
+}
+
+fn mempool_storage_basic_for_network(network: Network) -> Result<()> {
+    // Get transactions from the first 10 blocks of the Zcash blockchain
+    let unmined_transactions: Vec<_> = unmined_transactions_in_blocks(..=10, network).collect();
+
+    assert!(
+        MEMPOOL_TX_COUNT < unmined_transactions.len(),
+        "inconsistent MEMPOOL_TX_COUNT value for this test; decrease it"
+    );
+
+    // Use the sum of the costs of the first `MEMPOOL_TX_COUNT` transactions
+    // as the cost limit
+    let tx_cost_limit = unmined_transactions
+        .iter()
+        .take(MEMPOOL_TX_COUNT)
+        .map(|tx| tx.cost())
+        .sum();
+
+    // Create an empty storage
+    let mut storage: Storage = Storage::new(&config::Config {
+        tx_cost_limit,
+        ..Default::default()
+    });
+
+    // Insert them all to the storage
+    let mut maybe_inserted_transactions = Vec::new();
+    let mut some_rejected_transactions = Vec::new();
+    for unmined_transaction in unmined_transactions.clone() {
+        let result = storage.insert(unmined_transaction.clone());
+        match result {
+            Ok(_) => {
+                // While the transaction was inserted here, it can be rejected later.
+                maybe_inserted_transactions.push(unmined_transaction);
+            }
+            Err(_) => {
+                // Other transactions can be rejected on a successful insert,
+                // so not all rejected transactions will be added.
+                // Note that `some_rejected_transactions` can be empty since `insert` only
+                // returns a rejection error if the transaction being inserted is the one
+                // that was randomly evicted.
+                some_rejected_transactions.push(unmined_transaction);
+            }
+        }
+    }
+    // Since transactions are rejected randomly we can't test exact numbers.
+    // We know the first MEMPOOL_TX_COUNT must have been inserted successfully.
+    assert!(maybe_inserted_transactions.len() >= MEMPOOL_TX_COUNT);
+    assert_eq!(
+        some_rejected_transactions.len() + maybe_inserted_transactions.len(),
+        unmined_transactions.len()
+    );
+
+    // Test if the actual number of inserted/rejected transactions is consistent.
+    assert!(storage.verified.transaction_count() <= maybe_inserted_transactions.len());
+    assert!(storage.rejected_transaction_count() >= some_rejected_transactions.len());
+
+    // Test if rejected transactions were actually rejected.
+    for tx in some_rejected_transactions.iter() {
+        assert!(!storage.contains_transaction_exact(&tx.transaction.id));
+    }
+
+    // Query all the ids we have for rejected, get back `total - MEMPOOL_SIZE`
+    let all_ids: HashSet<UnminedTxId> = unmined_transactions
+        .iter()
+        .map(|tx| tx.transaction.id)
+        .collect();
+
+    // Convert response to a `HashSet`, because the order of the response doesn't matter.
+    let all_rejected_ids: HashSet<UnminedTxId> =
+        storage.rejected_transactions(all_ids).into_iter().collect();
+
+    let some_rejected_ids = some_rejected_transactions
+        .iter()
+        .map(|tx| tx.transaction.id)
+        .collect::<HashSet<_>>();
+
+    // Test if the rejected transactions we have are a subset of the actually
+    // rejected transactions.
+    assert!(some_rejected_ids.is_subset(&all_rejected_ids));
+
+    Ok(())
+}
+
 #[test]
 fn mempool_storage_crud_same_effects_mainnet() {
     zebra_test::init();

zebrad/src/components/mempool/tests/prop.rs

@@ -1,12 +1,13 @@
 //! Randomised property tests for the mempool.
 
+use proptest::collection::vec;
 use proptest::prelude::*;
 use proptest_derive::Arbitrary;
 
 use tokio::time;
 use tower::{buffer::Buffer, util::BoxService};
 
-use zebra_chain::{parameters::Network, transaction::VerifiedUnminedTx};
+use zebra_chain::{block, parameters::Network, transaction::VerifiedUnminedTx};
 use zebra_consensus::{error::TransactionError, transaction as tx};
 use zebra_network as zn;
 use zebra_state::{self as zs, ChainTipBlock, ChainTipSender};
@@ -26,6 +27,8 @@ type MockState = MockService<zs::Request, zs::Response, PropTestAssertion>;
 /// A [`MockService`] representing the Zebra transaction verifier service.
 type MockTxVerifier = MockService<tx::Request, tx::Response, PropTestAssertion, TransactionError>;
 
+const CHAIN_LENGTH: usize = 10;
+
 proptest! {
     /// Test if the mempool storage is cleared on a chain reset.
     #[test]
@@ -81,6 +84,94 @@ proptest! {
         })?;
     }
 
+    /// Test if the mempool storage is cleared on multiple chain resets.
+    #[test]
+    fn storage_is_cleared_on_chain_resets(
+        network in any::<Network>(),
+        mut previous_chain_tip in any::<ChainTipBlock>(),
+        mut transactions in vec(any::<VerifiedUnminedTx>(), 0..CHAIN_LENGTH),
+        fake_chain_tips in vec(any::<FakeChainTip>(), 0..CHAIN_LENGTH),
+    ) {
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .expect("Failed to create Tokio runtime");
+        let _guard = runtime.enter();
+
+        runtime.block_on(async move {
+            let (
+                mut mempool,
+                mut peer_set,
+                mut state_service,
+                mut tx_verifier,
+                mut recent_syncs,
+                mut chain_tip_sender,
+            ) = setup(network);
+
+            time::pause();
+
+            mempool.enable(&mut recent_syncs).await;
+
+            // Set the initial chain tip.
+            chain_tip_sender.set_best_non_finalized_tip(previous_chain_tip.clone());
+
+            // Call the mempool so that it is aware of the initial chain tip.
+            mempool.dummy_call().await;
+
+            for (fake_chain_tip, transaction) in fake_chain_tips.iter().zip(transactions.iter_mut()) {
+                // Obtain a new chain tip based on the previous one.
+                let chain_tip = fake_chain_tip.to_chain_tip_block(&previous_chain_tip);
+
+                // Adjust the transaction expiry height based on the new chain
+                // tip height so that the mempool does not evict the transaction
+                // when there is a chain growth.
+                if let Some(expiry_height) = transaction.transaction.transaction.expiry_height() {
+                    if chain_tip.height >= expiry_height {
+                        let mut tmp_tx = (*transaction.transaction.transaction).clone();
+
+                        // Set a new expiry height that is greater than the
+                        // height of the current chain tip.
+                        *tmp_tx.expiry_height_mut() = block::Height(chain_tip.height.0 + 1);
+                        transaction.transaction = tmp_tx.into();
+                    }
+                }
+
+                // Insert the dummy transaction into the mempool.
+                mempool
+                    .storage()
+                    .insert(transaction.clone())
+                    .expect("Inserting a transaction should succeed");
+
+                // Set the new chain tip.
+                chain_tip_sender.set_best_non_finalized_tip(chain_tip.clone());
+
+                // Call the mempool so that it is aware of the new chain tip.
+                mempool.dummy_call().await;
+
+                match fake_chain_tip {
+                    FakeChainTip::Grow(_) => {
+                        // The mempool should not be empty because we had a regular chain growth.
+                        prop_assert_ne!(mempool.storage().transaction_count(), 0);
+                    }
+
+                    FakeChainTip::Reset(_) => {
+                        // The mempool should be empty because we had a chain tip reset.
+                        prop_assert_eq!(mempool.storage().transaction_count(), 0);
+                    },
+                }
+
+                // Remember the current chain tip so that the next one can refer to it.
+                previous_chain_tip = chain_tip;
+            }
+
+            peer_set.expect_no_requests().await?;
+            state_service.expect_no_requests().await?;
+            tx_verifier.expect_no_requests().await?;
+
+            Ok(())
+        })?;
+    }
+
     /// Test if the mempool storage is cleared if the syncer falls behind and starts to catch up.
     #[test]
     fn storage_is_cleared_if_syncer_falls_behind(
@@ -185,3 +276,21 @@ enum FakeChainTip {
     Grow(ChainTipBlock),
     Reset(ChainTipBlock),
 }
+
+impl FakeChainTip {
+    /// Returns a new [`ChainTipBlock`] placed on top of the previous block if
+    /// the chain is supposed to grow. Otherwise returns a [`ChainTipBlock`]
+    /// that does not reference the previous one.
+    fn to_chain_tip_block(&self, previous: &ChainTipBlock) -> ChainTipBlock {
+        match self {
+            Self::Grow(chain_tip_block) => ChainTipBlock {
+                hash: chain_tip_block.hash,
+                height: block::Height(previous.height.0 + 1),
+                transaction_hashes: chain_tip_block.transaction_hashes.clone(),
+                previous_block_hash: previous.hash,
+            },
+
+            Self::Reset(chain_tip_block) => chain_tip_block.clone(),
+        }
+    }
+}