From 23705c5cecbde0cc680c9151d7409db147e70e8a Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Mon, 2 Dec 2024 14:16:28 -0600
Subject: [PATCH] tests

---
 .../thread_aware_account_locks.rs             | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/core/src/banking_stage/transaction_scheduler/thread_aware_account_locks.rs b/core/src/banking_stage/transaction_scheduler/thread_aware_account_locks.rs
index a0a5903a61d097..f3e3d3f8d683af 100644
--- a/core/src/banking_stage/transaction_scheduler/thread_aware_account_locks.rs
+++ b/core/src/banking_stage/transaction_scheduler/thread_aware_account_locks.rs
@@ -755,4 +755,39 @@ mod tests {
         let any_threads = ThreadSet::any(MAX_THREADS);
         assert_eq!(any_threads.num_threads(), MAX_THREADS as u32);
     }
+
+    #[test]
+    fn test_thread_set_iter() {
+        let mut thread_set = ThreadSet::none();
+        assert!(thread_set.contained_threads_iter().next().is_none());
+
+        thread_set.insert(4);
+        assert_eq!(
+            thread_set.contained_threads_iter().collect::<Vec<_>>(),
+            vec![4]
+        );
+
+        thread_set.insert(5);
+        assert_eq!(
+            thread_set.contained_threads_iter().collect::<Vec<_>>(),
+            vec![4, 5]
+        );
+        thread_set.insert(63);
+        assert_eq!(
+            thread_set.contained_threads_iter().collect::<Vec<_>>(),
+            vec![4, 5, 63]
+        );
+
+        thread_set.remove(5);
+        assert_eq!(
+            thread_set.contained_threads_iter().collect::<Vec<_>>(),
+            vec![4, 63]
+        );
+
+        let thread_set = ThreadSet::any(64);
+        assert_eq!(
+            thread_set.contained_threads_iter().collect::<Vec<_>>(),
+            (0..64).collect::<Vec<_>>()
+        );
+    }
 }