test: add unit test for PendingRequests

lambdaclass · Jan 25, 2025 · c31409c · c31409c
1 parent c9669b1
commit c31409c
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 2 deletions.
diff --git a/crates/networking/p2p/rlpx/tx_requests.rs b/crates/networking/p2p/rlpx/tx_requests.rs
diff --git a/crates/storage/store/pending_requests.rs b/crates/storage/store/pending_requests.rs
@@ -8,7 +8,7 @@ use ethrex_core::H512;
 
 const STALE_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct TransactionRequest {
     pub id: u64,
     pub transaction_hashes: Vec<H256>,
@@ -36,7 +36,7 @@ impl TransactionRequest {
         self.timestamp + STALE_REQUEST_TIMEOUT < time
     }
 }
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub(crate) struct PendingRequests {
     pending_transactions: HashSet<H256>,
     requests_by_node_id: HashMap<H512, Vec<TransactionRequest>>,
@@ -117,3 +117,89 @@ impl PendingRequests {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use ethereum_types::H256;
+    use ethrex_core::H512;
+    use std::thread::sleep;
+
+    fn get_transaction_request() -> (H512, Vec<H256>, TransactionRequest) {
+        let node_id = H512::random();
+        let tx_hashes = vec![H256::random(), H256::random()];
+        let request = TransactionRequest::new(1, tx_hashes.clone(), vec![], vec![]);
+        (node_id, tx_hashes, request)
+    }
+
+    #[test]
+    fn test_transaction_request_staleness() {
+        let request = TransactionRequest::new(1, vec![], vec![], vec![]);
+        sleep(STALE_REQUEST_TIMEOUT + Duration::from_secs(1));
+        assert!(request.is_stale(Instant::now()));
+    }
+
+    #[test]
+    fn test_store_request() {
+        let mut pending = PendingRequests::new();
+        let (node_id, tx_hashes, request) = get_transaction_request();
+
+        pending.store_request(node_id, request.clone());
+
+        assert!(pending.contains_txs(&tx_hashes[0]));
+        assert!(pending.contains_txs(&tx_hashes[1]));
+        assert_eq!(pending.requests_by_node_id[&node_id].len(), 1);
+        assert_eq!(pending.requests_by_node_id[&node_id][0], request);
+        assert_eq!(pending.get_pending_request(&node_id, 1).unwrap(), request);
+    }
+
+    #[test]
+    fn test_remove_pending_request() {
+        let mut pending = PendingRequests::new();
+        let (node_id, tx_hashes, request) = get_transaction_request();
+
+        pending.store_request(node_id, request);
+        pending.remove_pending_request(&node_id, 1);
+
+        assert!(!pending.contains_txs(&tx_hashes[0]));
+        assert!(!pending.contains_txs(&tx_hashes[1]));
+        assert!(pending.requests_by_node_id[&node_id].is_empty());
+    }
+
+    #[test]
+    fn test_remove_peer_requests() {
+        let mut pending = PendingRequests::new();
+        let (node_id, tx_hashes, request) = get_transaction_request();
+
+        pending.store_request(node_id, request);
+        pending.remove_peer_requests(&node_id);
+
+        assert!(!pending.contains_txs(&tx_hashes[0]));
+        assert!(!pending.contains_txs(&tx_hashes[1]));
+        assert!(!pending.requests_by_node_id.contains_key(&node_id));
+    }
+
+    #[test]
+    fn test_remove_stale_requests() {
+        let mut pending = PendingRequests::new();
+        let (node_id, stale_tx_hashes, stale_request) = get_transaction_request();
+
+        pending.store_request(node_id, stale_request);
+
+        sleep(STALE_REQUEST_TIMEOUT + Duration::from_secs(1));
+
+        let (_, fresh_tx_hashes, fresh_request) = get_transaction_request();
+
+        pending.store_request(node_id, fresh_request);
+
+        pending.remove_stale_requests(&node_id);
+
+        assert_eq!(pending.requests_by_node_id[&node_id].len(), 1);
+        assert_eq!(pending.requests_by_node_id[&node_id][0].id, 1);
+        assert!(!pending.contains_txs(&stale_tx_hashes[0]));
+        assert!(!pending.contains_txs(&stale_tx_hashes[1]));
+        assert!(pending.contains_txs(&fresh_tx_hashes[0]));
+        assert!(pending.contains_txs(&fresh_tx_hashes[1]));
+    }
+}