From f38b0b77158eb831ef8772fc6ed62c9a1938bac3 Mon Sep 17 00:00:00 2001
From: Ayelet Zilber <ayelet@starkware.co>
Date: Mon, 8 Jul 2024 14:34:29 +0300
Subject: [PATCH] test(mempool): add mempool state struct and new function

---
 crates/mempool/src/mempool_test.rs | 33 ++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs
index 05fcefb3..d94e73be 100644
--- a/crates/mempool/src/mempool_test.rs
+++ b/crates/mempool/src/mempool_test.rs
@@ -14,6 +14,31 @@ use crate::mempool::{Mempool, MempoolInput, TransactionReference};
 use crate::transaction_pool::TransactionPool;
 use crate::transaction_queue::TransactionQueue;
 
+/// `MempoolState` is a utility struct for testing mempool scenarios. It allows the creation of
+/// different mempool states, without following the regular mempool behavior.
+struct MempoolState {
+    tx_pool: TransactionPool,
+    tx_queue: TransactionQueue,
+}
+
+impl MempoolState {
+    fn new<T>(pool_txs: T, queue_txs: T) -> Self
+    where
+        T: IntoIterator<Item = ThinTransaction>,
+    {
+        let tx_pool = pool_txs.into_iter().collect();
+        let tx_queue = queue_txs.into_iter().collect();
+        MempoolState { tx_pool, tx_queue }
+    }
+}
+
+impl From<MempoolState> for Mempool {
+    fn from(mempool_state: MempoolState) -> Mempool {
+        let MempoolState { tx_pool, tx_queue } = mempool_state;
+        Mempool { tx_pool, tx_queue }
+    }
+}
+
 impl FromIterator<ThinTransaction> for TransactionPool {
     fn from_iter<T: IntoIterator<Item = ThinTransaction>>(txs: T) -> TransactionPool {
         let mut pool = TransactionPool::default();
@@ -98,13 +123,13 @@ fn test_get_txs(#[case] requested_txs: usize) {
     let input_tip_10_address_2 = add_tx_input!(tip: 10, tx_hash: 3, sender_address: "0x2");
 
     let txs = [
-        input_tip_50_address_0.clone(),
-        input_tip_100_address_1.clone(),
-        input_tip_10_address_2.clone(),
+        input_tip_50_address_0.clone().tx,
+        input_tip_100_address_1.clone().tx,
+        input_tip_10_address_2.clone().tx,
     ];
     let n_txs = txs.len();
 
-    let mut mempool = Mempool::new(txs).unwrap();
+    let mut mempool: Mempool = MempoolState::new(txs.clone(), txs).into();
 
     let sorted_txs =
         [input_tip_100_address_1.tx, input_tip_50_address_0.tx, input_tip_10_address_2.tx];