Added unit tests for Election Algorithm changes

Signed-off-by: Abhijit Gadgil <[email protected]>

components/butterfly/tests/common/mod.rs

@@ -368,6 +368,28 @@ impl SwimNet {
         }
     }
 
+    /// Partition a node from the rest of the network
+    pub fn partition_node(&self, idx: usize) {
+        println!("Partitioning {} from the network.", idx);
+        for i in 0..self.members.len() {
+            if i != idx {
+                self.block(idx, i);
+                self.block(i, idx);
+            }
+        }
+    }
+
+    /// UnPartition a node from the rest of the network
+    pub fn unpartition_node(&self, idx: usize) {
+        println!("UnPartitioning {} from the network.", idx);
+        for i in 0..self.members.len() {
+            if i != idx {
+                self.unblock(idx, i);
+                self.unblock(i, idx);
+            }
+        }
+    }
+
     /// # Locking (see locking.md)
     /// * `MemberList::entries` (read)
     pub fn wait_for_health_of_mlr(&self,

components/butterfly/tests/rumor/election.rs

@@ -1,3 +1,7 @@
+use std::iter::FromIterator;
+
+use rand::prelude::SliceRandom;
+
 use crate::btest;
 use habitat_butterfly::{member::Health,
                         rumor::{election::ElectionStatus,
@@ -92,8 +96,6 @@ fn five_members_elect_a_new_leader_when_the_old_one_dies() {
 #[test]
 #[allow(clippy::cognitive_complexity)]
 fn five_members_elect_a_new_leader_when_they_are_quorum_partitioned() {
-    env_logger::init();
-
     let mut net = btest::SwimNet::new_with_suitability_rhw(vec![1, 0, 0, 0, 0]);
     net[0].myself().lock_smw().set_persistent();
     net[4].myself().lock_smw().set_persistent();
@@ -173,3 +175,110 @@ fn five_members_elect_a_new_leader_when_they_are_quorum_partitioned() {
               assert_eq!(new_leader_id.as_ref(), Some(&e.member_id));
           });
 }
+
+#[test]
+#[allow(clippy::cognitive_complexity)]
+fn three_persistent_members_reelect_same_leader_follower_partition() {
+    // Server '0' will be the leader always.
+    let mut net = btest::SwimNet::new_with_suitability_rhw(vec![1, 0, 0]);
+    net[0].myself().lock_smw().set_persistent();
+    net[1].myself().lock_smw().set_persistent();
+    net[2].myself().lock_smw().set_persistent();
+
+    net.add_service(0, "core/foobar/1.2.3/20241010101010");
+    net.add_service(1, "core/foobar/1.2.3/20241010101010");
+    net.add_service(2, "core/foobar/1.2.3/20241010101010");
+    net.add_election(0, "foobar");
+    net.connect_smr(0, 1);
+    net.connect_smr(1, 2);
+    assert_wait_for_health_of_mlr!(net, [0..3, 0..3], Health::Alive);
+    assert_wait_for_election_status!(net, [0..3], "foobar.prod", ElectionStatus::Finished);
+    assert_wait_for_equal_election!(net, [0..3, 0..3], "foobar.prod");
+
+    let leader_id = net[0].election_store
+                          .lock_rsr()
+                          .service_group("foobar.prod")
+                          .map_rumor(Election::const_id(), |e| e.member_id.clone());
+
+    assert_eq!(leader_id, Some(net[0].member_id().to_string()));
+
+    net.partition(0..2, 2..3);
+    assert_wait_for_health_of_mlr!(net, [0..2, 2..3], Health::Confirmed);
+    assert_wait_for_election_status!(net, 0, "foobar.prod", ElectionStatus::Finished);
+    assert_wait_for_election_status!(net, 1, "foobar.prod", ElectionStatus::Finished);
+
+    net[2].restart_elections_rsw_mlr_rhw_msr(FeatureFlag::empty());
+    assert_wait_for_election_status!(net, 2, "foobar.prod", ElectionStatus::NoQuorum);
+
+    net.unpartition(0..2, 2..3);
+    assert_wait_for_health_of_mlr!(net, [0..2, 2..3], Health::Alive);
+    assert_wait_for_election_status!(net, 2, "foobar.prod", ElectionStatus::Finished);
+
+    let new_leader_id = net[2].election_store
+                              .lock_rsr()
+                              .service_group("foobar.prod")
+                              .map_rumor(Election::const_id(), |e| e.member_id.clone());
+
+    assert_eq!(leader_id, new_leader_id,
+               "OLD: {:?}, NEW: {:?}",
+               leader_id, new_leader_id);
+}
+
+#[test]
+#[allow(clippy::cognitive_complexity)]
+fn five_persistent_members_same_leader_multiple_non_quorum_partitions() {
+    let mut net = btest::SwimNet::new_with_suitability_rhw(vec![1, 0, 0, 0, 0]);
+    for i in 0..5 {
+        net[i].myself().lock_smw().set_persistent();
+        net.add_service(i, "core/foobar/1.2.3/20241010101010");
+    }
+    net.add_election(0, "foobar");
+
+    for servers in Vec::from_iter(0..5).windows(2) {
+        net.connect_smr(servers[0], servers[1])
+    }
+    assert_wait_for_health_of_mlr!(net, [0..5, 0..5], Health::Alive);
+    assert_wait_for_election_status!(net, [0..5], "foobar.prod", ElectionStatus::Finished);
+    assert_wait_for_equal_election!(net, [0..3, 0..3], "foobar.prod");
+
+    let leader_id = net[0].election_store
+                          .lock_rsr()
+                          .service_group("foobar.prod")
+                          .map_rumor(Election::const_id(), |e| e.member_id.clone());
+
+    assert_eq!(leader_id, Some(net[0].member_id().to_string()));
+
+    // Making sure - running multiple times after a subset of follower (non-quorum) is partitioned
+    // and reconnected the leader stays the same.
+    let mut rng = rand::thread_rng();
+    for _ in 0..2 {
+        let idxes = Vec::from_iter(1_usize..5_usize).choose_multiple(&mut rng, 2)
+                                                    .map(|x| *x)
+                                                    .collect::<Vec<usize>>();
+        for idx in idxes.iter() {
+            println!("idx: {}", idx);
+            net.partition_node(*idx);
+            assert_wait_for_health_of_mlr!(net, *idx, Health::Confirmed);
+        }
+        let first = idxes[0];
+        net[first].restart_elections_rsw_mlr_rhw_msr(FeatureFlag::empty());
+        assert_wait_for_election_status!(net, 0, "foobar.prod", ElectionStatus::Finished);
+
+        assert_wait_for_election_status!(net, first, "foobar.prod", ElectionStatus::NoQuorum);
+
+        for idx in idxes.iter() {
+            net.unpartition_node(*idx);
+            assert_wait_for_health_of_mlr!(net, *idx, Health::Alive);
+        }
+        assert_wait_for_election_status!(net, first, "foobar.prod", ElectionStatus::Finished);
+
+        let new_leader_id = net[first].election_store
+                                      .lock_rsr()
+                                      .service_group("foobar.prod")
+                                      .map_rumor(Election::const_id(), |e| e.member_id.clone());
+
+        assert_eq!(leader_id, new_leader_id,
+                   "OLD: {:?}, NEW: {:?}",
+                   leader_id, new_leader_id);
+    }
+}