Add unit test for testing cluster agent leader election

Author: zhuminyi

pkg/clusteragent/api/leader_handler_test.go

@@ -29,19 +29,25 @@ func (m *mockLeaderEngine) GetLeaderIP() (string, error) {
 	return m.leaderIP, nil
 }
 
-// fakeLeaderForwarder is a fake implementation of the forwarder for testing purposes
-type fakeLeaderForwarder struct{}
+// fakeLeaderForwarder is a fake implementation of the forwarder for testing purposes.
+// It tracks leader IP changes and forward calls for verifying leadership transition behavior.
+type fakeLeaderForwarder struct {
+	currentLeaderIP  string
+	leaderIPChanges  []string
+	forwardCallCount int
+}
 
-// SetLeaderIP does nothing
-func (f *fakeLeaderForwarder) SetLeaderIP(_ string) {}
+func (f *fakeLeaderForwarder) SetLeaderIP(ip string) {
+	f.currentLeaderIP = ip
+	f.leaderIPChanges = append(f.leaderIPChanges, ip)
+}
 
-// GetLeaderIP does nothing
 func (f *fakeLeaderForwarder) GetLeaderIP() string {
-	return ""
+	return f.currentLeaderIP
 }
 
-// Forward returns ok
 func (f *fakeLeaderForwarder) Forward(w http.ResponseWriter, _ *http.Request) {
+	f.forwardCallCount++
 	w.WriteHeader(http.StatusOK)
 }
 
@@ -92,3 +98,155 @@ func TestRejectOrForwardLeaderQuery_AsLeader(t *testing.T) {
 
 	assert.False(t, lph.rejectOrForwardLeaderQuery(rw, req))
 }
+
+// TestRejectOrForwardLeaderQuery_LeaderToFollowerTransition tests the behavior when
+// leadership changes from leader to follower between requests.
+func TestRejectOrForwardLeaderQuery_LeaderToFollowerTransition(t *testing.T) {
+	mockEngine := &mockLeaderEngine{
+		isLeader: true,
+		leaderIP: "1.1.1.1",
+	}
+	forwarder := &fakeLeaderForwarder{}
+
+	lph := &LeaderProxyHandler{
+		leaderElectionEnabled: true,
+		le:                    mockEngine,
+		leaderForwarder:       forwarder,
+	}
+
+	// First request: we are the leader, should handle locally
+	rw1 := httptest.NewRecorder()
+	req1 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.False(t, lph.rejectOrForwardLeaderQuery(rw1, req1), "Should handle locally as leader")
+	assert.Equal(t, 0, forwarder.forwardCallCount, "Should not forward when leader")
+
+	// Simulate leadership loss
+	mockEngine.isLeader = false
+	mockEngine.leaderIP = "2.2.2.2" // New leader IP
+
+	// Second request: we lost leadership, should forward to new leader
+	rw2 := httptest.NewRecorder()
+	req2 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.True(t, lph.rejectOrForwardLeaderQuery(rw2, req2), "Should forward as follower")
+	assert.Equal(t, 1, forwarder.forwardCallCount, "Should forward once")
+	assert.Equal(t, "2.2.2.2", forwarder.currentLeaderIP, "Should update to new leader IP")
+}
+
+// TestRejectOrForwardLeaderQuery_FollowerToLeaderTransition tests the behavior when
+// leadership changes from follower to leader between requests.
+func TestRejectOrForwardLeaderQuery_FollowerToLeaderTransition(t *testing.T) {
+	mockEngine := &mockLeaderEngine{
+		isLeader: false,
+		leaderIP: "1.1.1.1",
+	}
+	forwarder := &fakeLeaderForwarder{}
+
+	lph := &LeaderProxyHandler{
+		leaderElectionEnabled: true,
+		le:                    mockEngine,
+		leaderForwarder:       forwarder,
+	}
+
+	// First request: we are a follower, should forward
+	rw1 := httptest.NewRecorder()
+	req1 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.True(t, lph.rejectOrForwardLeaderQuery(rw1, req1), "Should forward as follower")
+	assert.Equal(t, 1, forwarder.forwardCallCount, "Should forward once")
+
+	// Simulate gaining leadership
+	mockEngine.isLeader = true
+
+	// Second request: we became the leader, should handle locally
+	rw2 := httptest.NewRecorder()
+	req2 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.False(t, lph.rejectOrForwardLeaderQuery(rw2, req2), "Should handle locally as new leader")
+	assert.Equal(t, 1, forwarder.forwardCallCount, "Should not forward additional requests")
+}
+
+// TestRejectOrForwardLeaderQuery_LeaderIPChange tests that the forwarder is updated
+// when the leader IP changes while we remain a follower.
+func TestRejectOrForwardLeaderQuery_LeaderIPChange(t *testing.T) {
+	mockEngine := &mockLeaderEngine{
+		isLeader: false,
+		leaderIP: "1.1.1.1",
+	}
+	forwarder := &fakeLeaderForwarder{
+		currentLeaderIP: "1.1.1.1", // Already knows old leader
+	}
+
+	lph := &LeaderProxyHandler{
+		leaderElectionEnabled: true,
+		le:                    mockEngine,
+		leaderForwarder:       forwarder,
+	}
+
+	// First request: forward to current leader
+	rw1 := httptest.NewRecorder()
+	req1 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.True(t, lph.rejectOrForwardLeaderQuery(rw1, req1))
+	assert.Equal(t, 1, forwarder.forwardCallCount)
+	// IP didn't change, so SetLeaderIP should not have been called
+	assert.Equal(t, 0, len(forwarder.leaderIPChanges), "Should not update IP when unchanged")
+
+	// Simulate leader failover - new leader elected
+	mockEngine.leaderIP = "2.2.2.2"
+
+	// Second request: should detect IP change and update forwarder
+	rw2 := httptest.NewRecorder()
+	req2 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.True(t, lph.rejectOrForwardLeaderQuery(rw2, req2))
+	assert.Equal(t, 2, forwarder.forwardCallCount)
+	assert.Equal(t, 1, len(forwarder.leaderIPChanges), "Should update IP once")
+	assert.Equal(t, "2.2.2.2", forwarder.currentLeaderIP, "Should have new leader IP")
+
+	// Third request: IP hasn't changed again
+	rw3 := httptest.NewRecorder()
+	req3 := httptest.NewRequest("GET", "http://example.com/foo", nil)
+	assert.True(t, lph.rejectOrForwardLeaderQuery(rw3, req3))
+	assert.Equal(t, 3, forwarder.forwardCallCount)
+	assert.Equal(t, 1, len(forwarder.leaderIPChanges), "Should not update IP when unchanged")
+}
+
+// TestRejectOrForwardLeaderQuery_MultipleLeaderChanges tests multiple leadership
+// transitions in sequence.
+func TestRejectOrForwardLeaderQuery_MultipleLeaderChanges(t *testing.T) {
+	mockEngine := &mockLeaderEngine{
+		isLeader: true,
+		leaderIP: "1.1.1.1",
+	}
+	forwarder := &fakeLeaderForwarder{}
+
+	lph := &LeaderProxyHandler{
+		leaderElectionEnabled: true,
+		le:                    mockEngine,
+		leaderForwarder:       forwarder,
+	}
+
+	// Sequence of leadership states to test
+	transitions := []struct {
+		isLeader       bool
+		leaderIP       string
+		expectForward  bool
+		expectedIPSets int
+	}{
+		{true, "1.1.1.1", false, 0}, // We are leader
+		{false, "2.2.2.2", true, 1}, // Lost leadership, forward to 2.2.2.2
+		{false, "2.2.2.2", true, 1}, // Still follower, same leader
+		{false, "3.3.3.3", true, 2}, // Still follower, new leader 3.3.3.3
+		{true, "1.1.1.1", false, 2}, // Regained leadership
+		{false, "4.4.4.4", true, 3}, // Lost again, new leader 4.4.4.4
+	}
+
+	for i, tr := range transitions {
+		mockEngine.isLeader = tr.isLeader
+		mockEngine.leaderIP = tr.leaderIP
+
+		rw := httptest.NewRecorder()
+		req := httptest.NewRequest("GET", "http://example.com/foo", nil)
+
+		result := lph.rejectOrForwardLeaderQuery(rw, req)
+		assert.Equal(t, tr.expectForward, result, "Transition %d: unexpected forward decision", i)
+		assert.Equal(t, tr.expectedIPSets, len(forwarder.leaderIPChanges),
+			"Transition %d: unexpected number of IP updates", i)
+	}
+}