diff --git a/internal_test.go b/internal_test.go index a6317196a..516ada823 100644 --- a/internal_test.go +++ b/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)) + }) + }) +}) diff --git a/osscluster.go b/osscluster.go index b018cc9e4..20180464e 100644 --- a/osscluster.go +++ b/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) }