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 +}