Add Routing Policies Comprehensive Test Suite and Fix multi keyed aggregation for different step

Author: ofekshenawa

command.go

@@ -201,6 +201,8 @@ type Cmder interface {
 	stringArg(int) string
 	firstKeyPos() int8
 	SetFirstKeyPos(int8)
+	stepCount() int8
+	SetStepCount(int8)
 
 	readTimeout() *time.Duration
 	readReply(rd *proto.Reader) error
@@ -303,6 +305,7 @@ type baseCmd struct {
 	args         []interface{}
 	err          error
 	keyPos       int8
+	_stepCount   int8
 	rawVal       interface{}
 	_readTimeout *time.Duration
 	cmdType      CmdType
@@ -361,6 +364,14 @@ func (cmd *baseCmd) SetFirstKeyPos(keyPos int8) {
 	cmd.keyPos = keyPos
 }
 
+func (cmd *baseCmd) stepCount() int8 {
+	return cmd._stepCount
+}
+
+func (cmd *baseCmd) SetStepCount(stepCount int8) {
+	cmd._stepCount = stepCount
+}
+
 func (cmd *baseCmd) SetErr(e error) {
 	cmd.err = e
 }
@@ -402,6 +413,7 @@ func (cmd *baseCmd) cloneBaseCmd() baseCmd {
 		args:         args,
 		err:          cmd.err,
 		keyPos:       cmd.keyPos,
+		_stepCount:   cmd._stepCount,
 		rawVal:       cmd.rawVal,
 		_readTimeout: readTimeout,
 		cmdType:      cmd.cmdType,

osscluster_router.go

@@ -34,6 +34,13 @@ func (c *ClusterClient) routeAndRun(ctx context.Context, cmd Cmder, node *cluste
 		policy = c.cmdInfoResolver.GetCommandPolicy(ctx, cmd)
 	}
 
+	// Set stepCount from cmdInfo if not already set
+	if cmd.stepCount() == 0 {
+		if cmdInfo := c.cmdInfo(ctx, cmd.Name()); cmdInfo != nil && cmdInfo.StepCount > 0 {
+			cmd.SetStepCount(cmdInfo.StepCount)
+		}
+	}
+
 	if policy == nil {
 		return c.executeDefault(ctx, cmd, policy, node)
 	}
@@ -104,6 +111,10 @@ func (c *ClusterClient) executeOnAllShards(ctx context.Context, cmd Cmder, polic
 func (c *ClusterClient) executeMultiShard(ctx context.Context, cmd Cmder, policy *routing.CommandPolicy) error {
 	args := cmd.Args()
 	firstKeyPos := int(cmdFirstKeyPos(cmd))
+	stepCount := int(cmd.stepCount())
+	if stepCount == 0 {
+		stepCount = 1 // Default to 1 if not set
+	}
 
 	if firstKeyPos == 0 || firstKeyPos >= len(args) {
 		return fmt.Errorf("redis: multi-shard command %s has no key arguments", cmd.Name())
@@ -113,14 +124,20 @@ func (c *ClusterClient) executeMultiShard(ctx context.Context, cmd Cmder, policy
 	slotMap := make(map[int][]string)
 	keyOrder := make([]string, 0)
 
-	for i := firstKeyPos; i < len(args); i++ {
+	for i := firstKeyPos; i < len(args); i += stepCount {
 		key, ok := args[i].(string)
 		if !ok {
 			return fmt.Errorf("redis: non-string key at position %d: %v", i, args[i])
 		}
 
 		slot := hashtag.Slot(key)
 		slotMap[slot] = append(slotMap[slot], key)
+		for j := 1; j < stepCount; j++ {
+			if i+j >= len(args) {
+				break
+			}
+			slotMap[slot] = append(slotMap[slot], args[i+j].(string))
+		}
 		keyOrder = append(keyOrder, key)
 	}
 
@@ -371,6 +388,7 @@ func (c *ClusterClient) aggregateMultiSlotResults(ctx context.Context, cmd Cmder
 			}
 		}
 
+		// TODO: return multiple errors by order when we will implement multiple errors returning
 		if result.err != nil {
 			firstErr = result.err
 		}