From 2640287e8e1bd74cff5bd7b580fb132799e65142 Mon Sep 17 00:00:00 2001
From: Jos Dehaes <jos.dehaes@gmail.com>
Date: Thu, 6 Jul 2023 12:46:48 +0200
Subject: [PATCH] feat: proptest for pending

---
 node/narwhal/src/helpers/pending.rs | 41 +++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/node/narwhal/src/helpers/pending.rs b/node/narwhal/src/helpers/pending.rs
index 7912b5a36c..f06a197212 100644
--- a/node/narwhal/src/helpers/pending.rs
+++ b/node/narwhal/src/helpers/pending.rs
@@ -108,6 +108,7 @@ mod tests {
         ledger::{coinbase::PuzzleCommitment, narwhal::TransmissionID},
         prelude::{Rng, TestRng},
     };
+    use test_strategy::{proptest, Arbitrary};
 
     type CurrentNetwork = snarkvm::prelude::Testnet3;
 
@@ -168,4 +169,44 @@ mod tests {
         // Check empty again.
         assert!(pending.is_empty());
     }
+
+    #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+    pub struct Item {
+        pub id: usize,
+    }
+
+    #[derive(Arbitrary, Clone, Debug)]
+    pub struct PendingInput {
+        #[strategy(1..5_000usize)]
+        pub count: usize,
+    }
+
+    impl PendingInput {
+        pub fn to_pending(&self) -> Pending<Item> {
+            let pending = Pending::<Item>::new();
+            for i in 0..self.count {
+                pending.insert(Item { id: i }, SocketAddr::from(([127, 0, 0, 1], i as u16)), None);
+            }
+            pending
+        }
+    }
+
+    #[proptest]
+    fn test_pending_proptest(input: PendingInput) {
+        println!("input: {:?}", input);
+        let pending = input.to_pending();
+        assert_eq!(pending.len(), input.count);
+        assert!(!pending.is_empty());
+        assert!(!pending.contains(Item { id: input.count + 1 }));
+        assert_eq!(pending.get(Item { id: input.count + 1 }), None);
+        assert!(!pending.remove(Item { id: input.count + 1 }));
+        for i in 0..input.count {
+            assert!(pending.contains(Item { id: i }));
+            let peer_ip = SocketAddr::from(([127, 0, 0, 1], i as u16));
+            assert!(pending.contains_peer(Item { id: i }, peer_ip));
+            assert_eq!(pending.get(Item { id: i }), Some(HashSet::from([peer_ip])));
+            assert!(pending.remove(Item { id: i }));
+        }
+        assert!(pending.is_empty());
+    }
 }