From a5823b1e78d4baa6f1982fb2dc798bbec2e35b2e Mon Sep 17 00:00:00 2001
From: Gilad Chase <gilad@starkware.com>
Date: Tue, 9 Jul 2024 21:06:31 +0200
Subject: [PATCH] feat: add tx_pool checker for mempool test

- Requires adding accessors to the private tx_pools; these should not be
used outside of our own tests.

- Didn't add a standalone tx_pool checker, can't think of a reason to
  test the pool without the queue (but the reverse is sensible
  somewhat).

- New tester will soon be used to test multi-nonce

commit-id:ba871431
---
 crates/mempool/src/mempool.rs          |  5 +++++
 crates/mempool/src/mempool_test.rs     | 17 +++++++++++++++++
 crates/mempool/src/transaction_pool.rs |  5 +++++
 3 files changed, 27 insertions(+)

diff --git a/crates/mempool/src/mempool.rs b/crates/mempool/src/mempool.rs
index 53f0e8947..cf87dbf31 100644
--- a/crates/mempool/src/mempool.rs
+++ b/crates/mempool/src/mempool.rs
@@ -85,6 +85,11 @@ impl Mempool {
 
         Ok(())
     }
+
+    #[cfg(test)]
+    pub(crate) fn _tx_pool(&self) -> &TransactionPool {
+        &self.tx_pool
+    }
 }
 
 /// Provides a lightweight representation of a transaction for mempool usage (e.g., excluding
diff --git a/crates/mempool/src/mempool_test.rs b/crates/mempool/src/mempool_test.rs
index cc83bfbb4..60140fa60 100644
--- a/crates/mempool/src/mempool_test.rs
+++ b/crates/mempool/src/mempool_test.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use assert_matches::assert_matches;
 use itertools::{enumerate, zip_eq};
 use pretty_assertions::assert_eq;
@@ -74,6 +76,21 @@ fn mempool() -> Mempool {
     Mempool::empty()
 }
 
+// TODO(Ayelet): replace with MempoolState checker.
+#[track_caller]
+fn _verify_mempool_state_eq(
+    mempool: &Mempool,
+    expected_txs: &[ThinTransaction],
+    expected_queue: &[ThinTransaction],
+) {
+    check_mempool_queue_eq(mempool, expected_queue);
+
+    let expected_txs: HashMap<_, _> =
+        expected_txs.iter().cloned().map(|tx| (tx.tx_hash, tx)).collect();
+
+    assert_eq!(mempool._tx_pool()._tx_pool(), &expected_txs);
+}
+
 // Asserts that the transactions in the mempool are in ascending order as per the expected
 // transactions.
 #[track_caller]
diff --git a/crates/mempool/src/transaction_pool.rs b/crates/mempool/src/transaction_pool.rs
index 9c7f5249e..bd1003651 100644
--- a/crates/mempool/src/transaction_pool.rs
+++ b/crates/mempool/src/transaction_pool.rs
@@ -70,6 +70,11 @@ impl TransactionPool {
     ) -> Option<&TransactionReference> {
         self.txs_by_account.get(address, nonce)
     }
+
+    #[cfg(test)]
+    pub(crate) fn _tx_pool(&self) -> &HashMap<TransactionHash, ThinTransaction> {
+        &self.tx_pool
+    }
 }
 
 #[derive(Debug, Default)]