From f9c2a7be612980072abe28af6fbe357e0c16d362 Mon Sep 17 00:00:00 2001
From: Sina Darbouy <sinadarbouy@gmail.com>
Date: Thu, 19 Dec 2024 17:46:45 +0100
Subject: [PATCH] feat(raft): Add comprehensive test coverage for FSM
 operations

Add test cases covering:
- Weighted round-robin operations (single and batch updates)
- Round-robin index management
- Invalid command handling
- FSM snapshot restoration
- Node shutdown scenarios
---
 raft/raft_test.go | 186 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 186 insertions(+)

diff --git a/raft/raft_test.go b/raft/raft_test.go
index fa912a47..a2a90849 100644
--- a/raft/raft_test.go
+++ b/raft/raft_test.go
@@ -270,3 +270,189 @@ func TestRaftLeadershipAndFollowers(t *testing.T) {
 	}
 	assert.Equal(t, 1, newLeaderCount, "Expected exactly one new leader after original leader shutdown")
 }
+
+func TestWeightedRROperations(t *testing.T) {
+	fsm := NewFSM()
+
+	// Test updating weighted round robin
+	cmd := Command{
+		Type: CommandUpdateWeightedRR,
+		Payload: WeightedRRPayload{
+			GroupName: "test-group",
+			ProxyName: "proxy1",
+			Weight: WeightedProxy{
+				CurrentWeight:   10,
+				EffectiveWeight: 20,
+			},
+		},
+	}
+
+	data, err := json.Marshal(cmd)
+	require.NoError(t, err)
+
+	// Apply the command
+	result := fsm.Apply(&raft.Log{Data: data})
+	assert.Nil(t, result)
+
+	// Test retrieving the weight
+	weight, exists := fsm.GetWeightedRRState("test-group", "proxy1")
+	assert.True(t, exists)
+	assert.Equal(t, WeightedProxy{CurrentWeight: 10, EffectiveWeight: 20}, weight)
+
+	// Test batch update
+	batchCmd := Command{
+		Type: CommandUpdateWeightedRRBatch,
+		Payload: WeightedRRBatchPayload{
+			GroupName: "test-group",
+			Updates: map[string]WeightedProxy{
+				"proxy1": {CurrentWeight: 15, EffectiveWeight: 25},
+				"proxy2": {CurrentWeight: 20, EffectiveWeight: 30},
+			},
+		},
+	}
+
+	data, err = json.Marshal(batchCmd)
+	require.NoError(t, err)
+
+	result = fsm.Apply(&raft.Log{Data: data})
+	assert.Nil(t, result)
+
+	// Verify batch updates
+	weight, exists = fsm.GetWeightedRRState("test-group", "proxy1")
+	assert.True(t, exists)
+	assert.Equal(t, WeightedProxy{CurrentWeight: 15, EffectiveWeight: 25}, weight)
+
+	weight, exists = fsm.GetWeightedRRState("test-group", "proxy2")
+	assert.True(t, exists)
+	assert.Equal(t, WeightedProxy{CurrentWeight: 20, EffectiveWeight: 30}, weight)
+}
+
+func TestRoundRobinOperations(t *testing.T) {
+	fsm := NewFSM()
+
+	// Test adding round robin index
+	cmd := Command{
+		Type: CommandAddRoundRobinNext,
+		Payload: RoundRobinPayload{
+			GroupName: "test-group",
+			NextIndex: 5,
+		},
+	}
+
+	data, err := json.Marshal(cmd)
+	require.NoError(t, err)
+
+	// Apply the command
+	result := fsm.Apply(&raft.Log{Data: data})
+	assert.Nil(t, result)
+
+	// Test retrieving the index
+	index := fsm.GetRoundRobinNext("test-group")
+	assert.Equal(t, uint32(5), index)
+
+	// Test non-existent group
+	index = fsm.GetRoundRobinNext("non-existent")
+	assert.Equal(t, uint32(0), index)
+}
+
+func TestFSMInvalidCommands(t *testing.T) {
+	fsm := NewFSM()
+
+	// Test invalid command type
+	cmd := Command{
+		Type:    "INVALID_COMMAND",
+		Payload: nil,
+	}
+
+	data, err := json.Marshal(cmd)
+	require.NoError(t, err)
+
+	result := fsm.Apply(&raft.Log{Data: data})
+	err, isError := result.(error)
+	assert.True(t, isError, "expected result to be an error")
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "unknown command type")
+
+	// Test invalid payload
+	cmd = Command{
+		Type:    CommandAddConsistentHashEntry,
+		Payload: "invalid payload",
+	}
+
+	data, err = json.Marshal(cmd)
+	require.NoError(t, err)
+
+	result = fsm.Apply(&raft.Log{Data: data})
+	err, isError = result.(error)
+	assert.True(t, isError, "expected result to be an error")
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to convert payload")
+}
+
+func TestFSMRestoreSnapshot(t *testing.T) {
+	fsm := NewFSM()
+
+	// Create test data
+	testData := struct {
+		HashToBlock     map[string]string                   `json:"hashToBlock"`
+		RoundRobin      map[string]uint32                   `json:"roundRobin"`
+		WeightedRRState map[string]map[string]WeightedProxy `json:"weightedRrState"`
+	}{
+		HashToBlock: map[string]string{"test-hash": "test-block"},
+		RoundRobin:  map[string]uint32{"test-group": 5},
+		WeightedRRState: map[string]map[string]WeightedProxy{
+			"test-group": {
+				"proxy1": {CurrentWeight: 10, EffectiveWeight: 20},
+			},
+		},
+	}
+
+	// Create a pipe to simulate snapshot reading
+	reader, writer := io.Pipe()
+
+	// Write test data in a goroutine
+	go func() {
+		encoder := json.NewEncoder(writer)
+		err := encoder.Encode(testData)
+		assert.NoError(t, err)
+		writer.Close()
+	}()
+
+	// Restore from snapshot
+	err := fsm.Restore(reader)
+	assert.NoError(t, err)
+
+	// Verify restored data
+	blockName, exists := fsm.GetProxyBlock("test-hash")
+	assert.True(t, exists)
+	assert.Equal(t, "test-block", blockName)
+
+	index := fsm.GetRoundRobinNext("test-group")
+	assert.Equal(t, uint32(5), index)
+
+	weight, exists := fsm.GetWeightedRRState("test-group", "proxy1")
+	assert.True(t, exists)
+	assert.Equal(t, WeightedProxy{CurrentWeight: 10, EffectiveWeight: 20}, weight)
+}
+
+func TestNodeShutdown(t *testing.T) {
+	logger := setupTestLogger()
+	tempDir := t.TempDir()
+
+	config := config.Raft{
+		NodeID:      "shutdown-test-node",
+		Address:     "127.0.0.1:0",
+		IsBootstrap: true,
+		Directory:   tempDir,
+	}
+
+	node, err := NewRaftNode(logger, config)
+	require.NoError(t, err)
+
+	// Test multiple shutdowns
+	err = node.Shutdown()
+	assert.NoError(t, err)
+
+	err = node.Shutdown()
+	assert.NoError(t, err) // Should not error on multiple shutdowns
+}