add test

Author: gregcusack

gossip/src/push_active_set.rs

@@ -53,6 +53,17 @@ impl Default for PushActiveSet {
     }
 }
 
+#[cfg(test)]
+impl PushActiveSet {
+    fn new_static() -> Self {
+        Self {
+            entries: Default::default(),
+            alpha: DEFAULT_ALPHA,
+            mode: WeightingMode::Static,
+        }
+    }
+}
+
 // Keys are gossip nodes to push messages to.
 // Values are which origins the node has pruned.
 #[derive(Default)]
@@ -148,7 +159,6 @@ impl PushActiveSet {
                 let f_scaled = ((num_unstaked * SCALE as usize) + nodes.len() / 2) / nodes.len();
                 let alpha_target = ALPHA_MIN + (f_scaled as u64).min(SCALE);
                 self.alpha = filter_alpha(self.alpha, alpha_target, ALPHA_MIN, ALPHA_MAX);
-                println!("alpha_target: {}, alpha: {}", alpha_target, self.alpha);
 
                 for (k, entry) in self.entries.iter_mut().enumerate() {
                     let weights: Vec<u64> = buckets
@@ -297,15 +307,15 @@ mod tests {
     }
 
     #[test]
-    fn test_push_active_set() {
+    fn test_push_active_set_static_weighting() {
         const CLUSTER_SIZE: usize = 117;
         let mut rng = ChaChaRng::from_seed([189u8; 32]);
         let pubkey = Pubkey::new_unique();
         let nodes: Vec<_> = repeat_with(Pubkey::new_unique).take(20).collect();
         let stakes = repeat_with(|| rng.gen_range(1..MAX_STAKE));
         let mut stakes: HashMap<_, _> = nodes.iter().copied().zip(stakes).collect();
         stakes.insert(pubkey, rng.gen_range(1..MAX_STAKE));
-        let mut active_set = PushActiveSet::default();
+        let mut active_set = PushActiveSet::new_static();
         assert!(active_set.entries.iter().all(|entry| entry.0.is_empty()));
         active_set.rotate(&mut rng, 5, CLUSTER_SIZE, &nodes, &stakes);
         assert!(active_set.entries.iter().all(|entry| entry.0.len() == 5));
@@ -317,28 +327,9 @@ mod tests {
         }
         let other = &nodes[5];
         let origin = &nodes[17];
-
-        // Debug: print expected nodes
-        // let expected_indices = [13, 5, 18, 16, 0];
-        // println!("Expected nodes:");
-        // for (i, &k) in expected_indices.iter().enumerate() {
-        //     println!("  [{}] nodes[{}] = {:?}", i, k, nodes[k]);
-        // }
-
-        // // Debug: print actual nodes returned by get_nodes
-        // let actual_nodes: Vec<_> = active_set
-        //     .get_nodes(&pubkey, origin, |_| false, &stakes)
-        //     .collect();
-        // println!("Actual nodes returned by get_nodes:");
-        // for (i, node) in actual_nodes.iter().enumerate() {
-        //     println!("  [{}] = {:?}", i, node);
-        // }
-
         assert!(active_set
             .get_nodes(&pubkey, origin, |_| false, &stakes)
             .eq([13, 5, 18, 16, 0].into_iter().map(|k| &nodes[k])));
-
-        ///////////////////////////
         assert!(active_set
             .get_nodes(&pubkey, other, |_| false, &stakes)
             .eq([13, 18, 16, 0].into_iter().map(|k| &nodes[k])));
@@ -371,6 +362,63 @@ mod tests {
             .eq([16, 7, 11].into_iter().map(|k| &nodes[k])));
     }
 
+    #[test]
+    fn test_push_active_set_dynamic_weighting() {
+        const CLUSTER_SIZE: usize = 117;
+        let mut rng = ChaChaRng::from_seed([14u8; 32]);
+        let pubkey = Pubkey::new_unique();
+        let nodes: Vec<_> = repeat_with(Pubkey::new_unique).take(20).collect();
+        let stakes = repeat_with(|| rng.gen_range(1..MAX_STAKE));
+        let mut stakes: HashMap<_, _> = nodes.iter().copied().zip(stakes).collect();
+        stakes.insert(pubkey, rng.gen_range(1..MAX_STAKE));
+        let mut active_set = PushActiveSet::default();
+        assert!(active_set.entries.iter().all(|entry| entry.0.is_empty()));
+        active_set.rotate(&mut rng, 5, CLUSTER_SIZE, &nodes, &stakes);
+        assert!(active_set.entries.iter().all(|entry| entry.0.len() == 5));
+        // Assert that for all entries, each filter already prunes the key.
+        for entry in &active_set.entries {
+            for (node, filter) in entry.0.iter() {
+                assert!(filter.contains(node));
+            }
+        }
+        let other = &nodes[6];
+        let origin = &nodes[17];
+        assert!(active_set
+            .get_nodes(&pubkey, origin, |_| false, &stakes)
+            .eq([7, 6, 2, 4, 12].into_iter().map(|k| &nodes[k])));
+        assert!(active_set
+            .get_nodes(&pubkey, other, |_| false, &stakes)
+            .eq([7, 2, 4, 12].into_iter().map(|k| &nodes[k])));
+
+        active_set.prune(&pubkey, &nodes[6], &[*origin], &stakes);
+        active_set.prune(&pubkey, &nodes[11], &[*origin], &stakes);
+        active_set.prune(&pubkey, &nodes[4], &[*origin], &stakes);
+        assert!(active_set
+            .get_nodes(&pubkey, origin, |_| false, &stakes)
+            .eq([7, 2, 12].into_iter().map(|k| &nodes[k])));
+        assert!(active_set
+            .get_nodes(&pubkey, other, |_| false, &stakes)
+            .eq([7, 2, 4, 12].into_iter().map(|k| &nodes[k])));
+        active_set.rotate(&mut rng, 7, CLUSTER_SIZE, &nodes, &stakes);
+        assert!(active_set.entries.iter().all(|entry| entry.0.len() == 7));
+        assert!(active_set
+            .get_nodes(&pubkey, origin, |_| false, &stakes)
+            .eq([2, 12, 16, 9, 14].into_iter().map(|k| &nodes[k])));
+        assert!(active_set
+            .get_nodes(&pubkey, other, |_| false, &stakes)
+            .eq([2, 4, 12, 16, 9, 14].into_iter().map(|k| &nodes[k])));
+        let origins = [*origin, *other];
+        active_set.prune(&pubkey, &nodes[2], &origins, &stakes);
+        active_set.prune(&pubkey, &nodes[12], &origins, &stakes);
+        active_set.prune(&pubkey, &nodes[9], &origins, &stakes);
+        assert!(active_set
+            .get_nodes(&pubkey, origin, |_| false, &stakes)
+            .eq([16, 14].into_iter().map(|k| &nodes[k])));
+        assert!(active_set
+            .get_nodes(&pubkey, other, |_| false, &stakes)
+            .eq([4, 16, 14].into_iter().map(|k| &nodes[k])));
+    }
+
     #[test]
     fn test_push_active_set_entry() {
         const NUM_BLOOM_FILTER_ITEMS: usize = 100;