Fake timer for tests?

tinyraft.go

@@ -60,6 +60,7 @@ type RaftConfig struct {
     HeartbeatTimeout  time.Duration
     LeadershipTimeout time.Duration
     SendFunction      func(to NodeID, msg RaftMsg)
+    FakeTimer         func(forHeartbeat bool, duration time.Duration) <-chan time.Time
 }
 
 type Raft struct {
@@ -67,6 +68,7 @@ type Raft struct {
     config   RaftConfig
     send     func(to NodeID, msg RaftMsg)
     onChange func(state NodeState, term uint64, leader NodeID, followers []NodeID)
+    timerFn  func(forHeartbeat bool, duration time.Duration) <-chan time.Time
 
     mu sync.Mutex
 
@@ -89,10 +91,17 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
     if len(config.Nodes) == 0 {
         config.Nodes = append(config.Nodes, config.NodeID)
     }
+    timerFn := config.FakeTimer
+    if timerFn == nil {
+        timerFn = func(forHeartbeat bool, duration time.Duration) <-chan time.Time {
+            return time.After(duration)
+        }
+    }
     r := &Raft{
         ctx:     ctx,
         config:  config,
         send:    config.SendFunction,
+        timerFn: timerFn,
         term:    config.InitialTerm,
     }
     r.electionTimerCh = make(chan time.Duration, 1)
@@ -111,7 +120,7 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
                 case <-r.ctx.Done():
                     return
                 case dur, ok = <-r.electionTimerCh:
-                case <-time.After(dur):
+                case <-r.timerFn(false, dur):
                     r.Start()
                 }
             }
@@ -122,7 +131,7 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
             select {
             case <-r.ctx.Done():
                 return
-            case <-time.After(r.config.HeartbeatTimeout):
+            case <-r.timerFn(true, r.config.HeartbeatTimeout):
                 r.heartbeat()
             }
         }