From df6e1c9add9a3c84001eae6d784251599a69ae95 Mon Sep 17 00:00:00 2001
From: Eitan Seri-Levi <eserilev@ucsc.edu>
Date: Thu, 30 May 2024 16:09:22 +0200
Subject: [PATCH] Fix BitVectors TestRandom implementation (#5854)

* fix bitvector test random impl
---
 consensus/types/src/test_utils/test_random/bitfield.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/consensus/types/src/test_utils/test_random/bitfield.rs b/consensus/types/src/test_utils/test_random/bitfield.rs
index f73f7c18c5a..35176d389d0 100644
--- a/consensus/types/src/test_utils/test_random/bitfield.rs
+++ b/consensus/types/src/test_utils/test_random/bitfield.rs
@@ -26,6 +26,15 @@ impl<N: Unsigned + Clone> TestRandom for BitVector<N> {
     fn random_for_test(rng: &mut impl RngCore) -> Self {
         let mut raw_bytes = smallvec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)];
         rng.fill_bytes(&mut raw_bytes);
+        // If N isn't divisible by 8
+        // zero out bits greater than N
+        if let Some(last_byte) = raw_bytes.last_mut() {
+            let mut mask = 0;
+            for i in 0..N::to_usize() % 8 {
+                mask |= 1 << i;
+            }
+            *last_byte &= mask;
+        }
         Self::from_bytes(raw_bytes).expect("we generate a valid BitVector")
     }
 }