Merge pull request #441 from redis/cmd-set-slot

feat: SetSlot to override the calculated key slot of a command

internal/cmds/cmds.go

@@ -145,6 +145,16 @@ func (c *Completed) Slot() uint16 {
 	return c.ks
 }
 
+// SetSlot returns a new completed command with its key slot be overridden
+func (c Completed) SetSlot(key string) Completed {
+	if c.ks&NoSlot == NoSlot {
+		c.ks = NoSlot | slot(key)
+	} else {
+		c.ks = slot(key)
+	}
+	return c
+}
+
 // Cacheable represents a completed Redis command which supports server-assisted client side caching,
 // and it should be created by the Cache() of command builder.
 type Cacheable Completed

internal/cmds/cmds_test.go

@@ -160,16 +160,23 @@ func TestCompleted_CommandSlice(t *testing.T) {
 }
 
 func TestCompleted_Slots(t *testing.T) {
-	builder := NewBuilder(InitSlot)
-	c1 := builder.Get().Key("a").Build()
-	c2 := builder.Get().Key("b").Build()
-	c3 := Cacheable(c1)
-	c4 := Cacheable(c2)
-	if c1.Slot() == c2.Slot() {
-		t.Fatalf("unexpected same slot")
-	}
-	if c3.Slot() == c4.Slot() {
-		t.Fatalf("unexpected same slot")
+	for _, init := range []uint16{InitSlot, NoSlot} {
+		builder := NewBuilder(init)
+		c1 := builder.Get().Key("a").Build()
+		c2 := builder.Get().Key("b").Build()
+		c3 := Cacheable(c1)
+		c4 := Cacheable(c2)
+		c5 := c1.SetSlot("c")
+		c6 := c2.SetSlot("c")
+		if c1.Slot() == c2.Slot() {
+			t.Fatalf("unexpected same slot")
+		}
+		if c3.Slot() == c4.Slot() {
+			t.Fatalf("unexpected same slot")
+		}
+		if c5.Slot() != c6.Slot() {
+			t.Fatalf("unexpected different slot")
+		}
 	}
 }