From 43beef256f55fef94e2ca0aadf8c9f2dbcfbc3f4 Mon Sep 17 00:00:00 2001
From: Rostyslav Tyshko <tyshko.rostyslav@gmail.com>
Date: Wed, 16 Oct 2024 04:32:38 +0200
Subject: [PATCH] add test_drain_size

---
 mempool/src/lib.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/mempool/src/lib.rs b/mempool/src/lib.rs
index ce9f7e2..28b36b7 100644
--- a/mempool/src/lib.rs
+++ b/mempool/src/lib.rs
@@ -137,6 +137,20 @@ mod tests {
         assert_eq!(pool.len(), 1);
     }
 
+    #[test]
+    fn test_drain_size() {
+        let mut pool = MemPool::new();
+        pool.push_item(test_item_with_id(1));
+        pool.push_item(test_item_with_id(2));
+        pool.push_item(test_item_with_id(3));
+        pool.push_item(test_item_with_id(4));
+
+        let items = pool.drain_size(2);
+        assert_eq!(items, vec![test_item_with_id(1), test_item_with_id(2)]);
+        assert_eq!(pool.len(), 2);
+    }
+
+
 
     #[test]
     fn test_push_pop() {