From d11e4a0dc2dae8493129768259dcab35d6edefb7 Mon Sep 17 00:00:00 2001
From: jayice <1185430411@qq.com>
Date: Sat, 1 Jun 2024 01:53:25 +0800
Subject: [PATCH] support diasble election in ut

---
 test/test_node.cpp | 15 ++++++++++++---
 test/util.h        | 12 ++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/test/test_node.cpp b/test/test_node.cpp
index 407e16af..9638b621 100644
--- a/test/test_node.cpp
+++ b/test/test_node.cpp
@@ -454,14 +454,19 @@ TEST_P(NodeTest, LeaderFailWithWitness) {
     }
     cond.wait();
 
+    std::vector<braft::Node*> follower_nodes;
+    cluster.followers(&follower_nodes);
+    ASSERT_EQ(2, follower_nodes.size());
+    for (auto node : follower_nodes){
+        cluster.disable_election(node);
+    }
+
     // stop leader
     butil::EndPoint old_leader = leader->node_id().peer_id.addr;
     LOG(WARNING) << "stop leader " << leader->node_id();
     cluster.stop(leader->node_id().peer_id.addr);
 
     // apply something when follower
-    std::vector<braft::Node*> nodes;
-    cluster.followers(&nodes);
     cond.reset(10);
     for (int i = 0; i < 10; i++) {
         butil::IOBuf data;
@@ -472,10 +477,14 @@ TEST_P(NodeTest, LeaderFailWithWitness) {
         task.data = &data;
         task.done = NEW_APPLYCLOSURE(&cond, -1);
         // node 0 is witness;
-        nodes[1]->apply(task);
+        follower_nodes[1]->apply(task);
     }
     cond.wait();
 
+    for (auto node : follower_nodes){
+        cluster.enable_election(node);
+    }
+
     // elect new leader
     cluster.wait_leader();
     leader = cluster.leader();
diff --git a/test/util.h b/test/util.h
index 5cbcf241..7df358dd 100644
--- a/test/util.h
+++ b/test/util.h
@@ -239,6 +239,7 @@ class Cluster {
         int64_t throttle_throughput_bytes = 10 * 1024 * 1024;
         int64_t check_cycle = 10;
         _throttle = new braft::ThroughputSnapshotThrottle(throttle_throughput_bytes, check_cycle);
+        _infinite_election_timeout_ms = 1 << 31;
     }
     ~Cluster() {
         stop_all();
@@ -496,6 +497,16 @@ class Cluster {
         goto CHECK;
     }
 
+    void disable_election(braft::Node* node){
+        std::lock_guard<raft_mutex_t> guard(_mutex);
+        node->reset_election_timeout_ms(_infinite_election_timeout_ms, _max_clock_drift_ms);
+    }
+
+    void enable_election(braft::Node* node) {
+        std::lock_guard<raft_mutex_t> guard(_mutex);
+        node->reset_election_timeout_ms(_election_timeout_ms, _max_clock_drift_ms);
+    }
+
 private:
     void all_nodes(std::vector<butil::EndPoint>* addrs) {
         addrs->clear();
@@ -540,6 +551,7 @@ class Cluster {
     std::map<butil::EndPoint, brpc::Server*> _server_map;
     int32_t _election_timeout_ms;
     int32_t _max_clock_drift_ms;
+    int32_t _infinite_election_timeout_ms;
     raft_mutex_t _mutex;
     braft::SnapshotThrottle* _throttle;
 };