From 4cde14ef21c3dedda78b9d94e45a08b3259c083d Mon Sep 17 00:00:00 2001
From: Mohammad Nassar <mohammad@starkware.co>
Date: Sun, 15 Sep 2024 10:25:06 +0300
Subject: [PATCH] chore(mempool): add transaction queue content

---
 crates/mempool/src/transaction_queue.rs       |  4 +
 .../src/transaction_queue_test_utils.rs       | 93 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 crates/mempool/src/transaction_queue_test_utils.rs

diff --git a/crates/mempool/src/transaction_queue.rs b/crates/mempool/src/transaction_queue.rs
index 645b3a04a6..22f14045cf 100644
--- a/crates/mempool/src/transaction_queue.rs
+++ b/crates/mempool/src/transaction_queue.rs
@@ -12,6 +12,10 @@ use starknet_api::transaction::{
 
 use crate::mempool::TransactionReference;
 
+#[cfg(test)]
+#[path = "transaction_queue_test_utils.rs"]
+pub mod transaction_queue_test_utils;
+
 // A queue holding the transaction that with nonces that match account nonces.
 // Note: the derived comparison functionality considers the order guaranteed by the data structures
 // used.
diff --git a/crates/mempool/src/transaction_queue_test_utils.rs b/crates/mempool/src/transaction_queue_test_utils.rs
new file mode 100644
index 0000000000..9ae7c5cdc5
--- /dev/null
+++ b/crates/mempool/src/transaction_queue_test_utils.rs
@@ -0,0 +1,93 @@
+use std::collections::{BTreeSet, HashMap};
+
+use pretty_assertions::assert_eq;
+use starknet_api::core::ContractAddress;
+
+use crate::mempool::TransactionReference;
+use crate::transaction_queue::{PendingTransaction, PriorityTransaction, TransactionQueue};
+
+// Utils.
+
+/// Represents the internal content of the transaction queue.
+/// Enables customized (and potentially inconsistent) creation for unit testing.
+#[derive(Debug, Default)]
+pub struct TransactionQueueContent {
+    priority_queue: Option<BTreeSet<PriorityTransaction>>,
+    pending_queue: Option<BTreeSet<PendingTransaction>>,
+    address_to_tx: Option<HashMap<ContractAddress, TransactionReference>>,
+}
+
+impl TransactionQueueContent {
+    pub fn _assert_eq_priority_and_pending_queues(&self, tx_queue: &TransactionQueue) {
+        self._assert_eq_priority_queue(tx_queue);
+        self._assert_eq_pending_queue(tx_queue);
+    }
+
+    pub fn _assert_eq_priority_queue(&self, tx_queue: &TransactionQueue) {
+        assert_eq!(self.priority_queue.as_ref().unwrap(), &tx_queue.priority_queue);
+        assert_eq!(self.address_to_tx.as_ref().unwrap(), &tx_queue.address_to_tx);
+    }
+
+    pub fn _assert_eq_pending_queue(&self, tx_queue: &TransactionQueue) {
+        assert_eq!(self.pending_queue.as_ref().unwrap(), &tx_queue.pending_queue);
+        assert_eq!(self.address_to_tx.as_ref().unwrap(), &tx_queue.address_to_tx);
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct TransactionQueueContentBuilder {
+    _priority_queue: Option<BTreeSet<PriorityTransaction>>,
+    _pending_queue: Option<BTreeSet<PendingTransaction>>,
+    _address_to_tx: Option<HashMap<ContractAddress, TransactionReference>>,
+}
+
+impl TransactionQueueContentBuilder {
+    pub fn _new() -> Self {
+        Self::default()
+    }
+
+    pub fn _with_priority<P>(mut self, priority_txs: P) -> Self
+    where
+        P: IntoIterator<Item = TransactionReference>,
+    {
+        self._priority_queue = Some(priority_txs.into_iter().map(Into::into).collect());
+        self._address_to_tx.get_or_insert_with(HashMap::new).extend(
+            self._priority_queue.as_ref().unwrap().iter().map(|tx| (tx.sender_address, tx.0)),
+        );
+
+        self
+    }
+
+    pub fn _with_pending<P>(mut self, pending_txs: P) -> Self
+    where
+        P: IntoIterator<Item = TransactionReference>,
+    {
+        self._pending_queue = Some(pending_txs.into_iter().map(Into::into).collect());
+        self._address_to_tx.get_or_insert_with(HashMap::new).extend(
+            self._pending_queue.as_ref().unwrap().iter().map(|tx| (tx.sender_address, tx.0)),
+        );
+
+        self
+    }
+
+    pub fn _build(self) -> TransactionQueueContent {
+        TransactionQueueContent {
+            priority_queue: self._priority_queue,
+            pending_queue: self._pending_queue,
+            address_to_tx: self._address_to_tx,
+        }
+    }
+}
+
+impl From<TransactionQueueContent> for TransactionQueue {
+    fn from(tx_queue_content: TransactionQueueContent) -> Self {
+        let TransactionQueueContent { priority_queue, pending_queue, address_to_tx } =
+            tx_queue_content;
+        TransactionQueue {
+            priority_queue: priority_queue.unwrap_or_default(),
+            pending_queue: pending_queue.unwrap_or_default(),
+            address_to_tx: address_to_tx.unwrap_or_default(),
+            ..Default::default()
+        }
+    }
+}