From ab356e3c149826c1a3627dacd3089d6c71ade6ff Mon Sep 17 00:00:00 2001
From: Nan Yan <625518543@qq.com>
Date: Mon, 13 May 2024 13:43:08 +0000
Subject: [PATCH] test: add case for LeaderChangeHandler

---
 fastflow_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/fastflow_test.go b/fastflow_test.go
index ac9fe1f..a79a552 100644
--- a/fastflow_test.go
+++ b/fastflow_test.go
@@ -1,11 +1,13 @@
 package fastflow
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
 
 	"github.com/shiningrush/fastflow/pkg/entity"
+	"github.com/shiningrush/fastflow/pkg/event"
 	"github.com/shiningrush/fastflow/pkg/mod"
 	"github.com/shiningrush/fastflow/pkg/utils"
 	"github.com/stretchr/testify/assert"
@@ -190,3 +192,48 @@ tasks:
 		})
 	}
 }
+
+func Test_LeaderChangeHandler(t *testing.T) {
+	tests := []struct {
+		isLeader      bool
+		calledEnsured []bool
+		closerLenth   int
+	}{
+		{
+			isLeader:      true,
+			calledEnsured: []bool{},
+			closerLenth:   1,
+		},
+		{
+			isLeader:      false,
+			calledEnsured: []bool{true},
+			closerLenth:   0,
+		},
+	}
+
+	for _, tc := range tests {
+		called := []bool{}
+		keeper := &mod.MockKeeper{}
+		handler := &LeaderChangedHandler{
+			opt: &InitialOption{
+				Keeper: keeper,
+			},
+			leaderCloser: []mod.Closer{},
+		}
+
+		closer := &mod.MockCloser{}
+		closer.On("Close").Run(func(args mock.Arguments) {
+			called = append(called, true)
+		}).Return(nil)
+		handler.leaderCloser = append(handler.leaderCloser, closer)
+
+		mockEvent := &event.LeaderChanged{IsLeader: tc.isLeader, WorkerKey: "worker-key"}
+		ctx := context.Background()
+
+		// Simulate the event that leads to leadership
+		handler.Handle(ctx, mockEvent)
+
+		assert.Equal(t, tc.calledEnsured, called)
+		assert.Equal(t, tc.closerLenth, len(handler.leaderCloser))
+	}
+}