From c75adbc9c70cd22bd19ba75caff4535b58f2c5b4 Mon Sep 17 00:00:00 2001
From: Arni Hod <arnon@starkware.co>
Date: Tue, 7 Jan 2025 09:15:30 +0200
Subject: [PATCH] test(starknet_mempool): mempool content builder with state

---
 crates/starknet_mempool/src/mempool_test.rs | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/crates/starknet_mempool/src/mempool_test.rs b/crates/starknet_mempool/src/mempool_test.rs
index 95fd7d9d266..646fb321ac6 100644
--- a/crates/starknet_mempool/src/mempool_test.rs
+++ b/crates/starknet_mempool/src/mempool_test.rs
@@ -19,7 +19,7 @@ use starknet_mempool_types::errors::MempoolError;
 use starknet_mempool_types::mempool_types::AddTransactionArgs;
 
 use crate::communication::MempoolCommunicationWrapper;
-use crate::mempool::{Mempool, MempoolConfig, TransactionReference};
+use crate::mempool::{Mempool, MempoolConfig, MempoolState, TransactionReference};
 use crate::test_utils::{add_tx, add_tx_expect_error, commit_block, get_txs_and_assert_expected};
 use crate::transaction_pool::TransactionPool;
 use crate::transaction_queue::transaction_queue_test_utils::{
@@ -37,6 +37,7 @@ struct MempoolContent {
     config: MempoolConfig,
     tx_pool: Option<TransactionPool>,
     tx_queue_content: Option<TransactionQueueContent>,
+    state: Option<MempoolState>,
 }
 
 impl MempoolContent {
@@ -54,15 +55,14 @@ impl MempoolContent {
 
 impl From<MempoolContent> for Mempool {
     fn from(mempool_content: MempoolContent) -> Mempool {
-        let MempoolContent { tx_pool, tx_queue_content, config } = mempool_content;
+        let MempoolContent { config, tx_pool, tx_queue_content, state } = mempool_content;
         Mempool {
             config,
             tx_pool: tx_pool.unwrap_or_default(),
             tx_queue: tx_queue_content
                 .map(|content| content.complete_to_tx_queue())
                 .unwrap_or_default(),
-            // TODO: Add implementation when needed.
-            state: Default::default(),
+            state: state.unwrap_or_default(),
         }
     }
 }
@@ -72,6 +72,7 @@ struct MempoolContentBuilder {
     config: MempoolConfig,
     tx_pool: Option<TransactionPool>,
     tx_queue_content_builder: TransactionQueueContentBuilder,
+    state: Option<MempoolState>,
 }
 
 impl MempoolContentBuilder {
@@ -80,6 +81,7 @@ impl MempoolContentBuilder {
             config: MempoolConfig { enable_fee_escalation: false, ..Default::default() },
             tx_pool: None,
             tx_queue_content_builder: Default::default(),
+            state: None,
         }
     }
 
@@ -91,6 +93,12 @@ impl MempoolContentBuilder {
         self
     }
 
+    #[allow(dead_code)]
+    fn with_state(mut self, state: MempoolState) -> Self {
+        self.state = Some(state);
+        self
+    }
+
     fn with_priority_queue<Q>(mut self, queue_txs: Q) -> Self
     where
         Q: IntoIterator<Item = TransactionReference>,
@@ -123,6 +131,7 @@ impl MempoolContentBuilder {
             config: self.config,
             tx_pool: self.tx_pool,
             tx_queue_content: self.tx_queue_content_builder.build(),
+            state: self.state,
         }
     }