|
1 | 1 | package fastflow
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "context" |
4 | 5 | "fmt"
|
5 | 6 | "testing"
|
6 | 7 | "time"
|
7 | 8 |
|
8 | 9 | "github.com/shiningrush/fastflow/pkg/entity"
|
| 10 | + "github.com/shiningrush/fastflow/pkg/event" |
9 | 11 | "github.com/shiningrush/fastflow/pkg/mod"
|
10 | 12 | "github.com/shiningrush/fastflow/pkg/utils"
|
11 | 13 | "github.com/stretchr/testify/assert"
|
@@ -190,3 +192,48 @@ tasks:
|
190 | 192 | })
|
191 | 193 | }
|
192 | 194 | }
|
| 195 | + |
| 196 | +func Test_LeaderChangeHandler(t *testing.T) { |
| 197 | + tests := []struct { |
| 198 | + isLeader bool |
| 199 | + calledEnsured []bool |
| 200 | + closerLenth int |
| 201 | + }{ |
| 202 | + { |
| 203 | + isLeader: true, |
| 204 | + calledEnsured: []bool{}, |
| 205 | + closerLenth: 1, |
| 206 | + }, |
| 207 | + { |
| 208 | + isLeader: false, |
| 209 | + calledEnsured: []bool{true}, |
| 210 | + closerLenth: 0, |
| 211 | + }, |
| 212 | + } |
| 213 | + |
| 214 | + for _, tc := range tests { |
| 215 | + called := []bool{} |
| 216 | + keeper := &mod.MockKeeper{} |
| 217 | + handler := &LeaderChangedHandler{ |
| 218 | + opt: &InitialOption{ |
| 219 | + Keeper: keeper, |
| 220 | + }, |
| 221 | + leaderCloser: []mod.Closer{}, |
| 222 | + } |
| 223 | + |
| 224 | + closer := &mod.MockCloser{} |
| 225 | + closer.On("Close").Run(func(args mock.Arguments) { |
| 226 | + called = append(called, true) |
| 227 | + }).Return(nil) |
| 228 | + handler.leaderCloser = append(handler.leaderCloser, closer) |
| 229 | + |
| 230 | + mockEvent := &event.LeaderChanged{IsLeader: tc.isLeader, WorkerKey: "worker-key"} |
| 231 | + ctx := context.Background() |
| 232 | + |
| 233 | + // Simulate the event that leads to leadership |
| 234 | + handler.Handle(ctx, mockEvent) |
| 235 | + |
| 236 | + assert.Equal(t, tc.calledEnsured, called) |
| 237 | + assert.Equal(t, tc.closerLenth, len(handler.leaderCloser)) |
| 238 | + } |
| 239 | +} |
0 commit comments