Use correct slot for COUNTKEYSINSLOT command (#3327)

Author: khasanovbi

internal_test.go

@@ -352,3 +352,27 @@ var _ = Describe("withConn", func() {
 		Expect(client.connPool.Len()).To(Equal(1))
 	})
 })
+
+var _ = Describe("ClusterClient", func() {
+	var client *ClusterClient
+
+	BeforeEach(func() {
+		client = &ClusterClient{}
+	})
+
+	Describe("cmdSlot", func() {
+		It("select slot from args for GETKEYSINSLOT command", func() {
+			cmd := NewStringSliceCmd(ctx, "cluster", "getkeysinslot", 100, 200)
+
+			slot := client.cmdSlot(context.Background(), cmd)
+			Expect(slot).To(Equal(100))
+		})
+
+		It("select slot from args for COUNTKEYSINSLOT command", func() {
+			cmd := NewStringSliceCmd(ctx, "cluster", "countkeysinslot", 100)
+
+			slot := client.cmdSlot(context.Background(), cmd)
+			Expect(slot).To(Equal(100))
+		})
+	})
+})

osscluster.go

@@ -1856,7 +1856,7 @@ func (c *ClusterClient) cmdInfo(ctx context.Context, name string) *CommandInfo {
 
 func (c *ClusterClient) cmdSlot(ctx context.Context, cmd Cmder) int {
 	args := cmd.Args()
-	if args[0] == "cluster" && args[1] == "getkeysinslot" {
+	if args[0] == "cluster" && (args[1] == "getkeysinslot" || args[1] == "countkeysinslot") {
 		return args[2].(int)
 	}