@@ -782,3 +782,82 @@ var _ = Describe("Ring Tx timeout", func() {
782782 testTimeout ()
783783 })
784784})
785+
786+ var _ = Describe ("Ring GetShardClients and GetShardClientForKey" , func () {
787+ var ring * redis.Ring
788+
789+ BeforeEach (func () {
790+ ring = redis .NewRing (& redis.RingOptions {
791+ Addrs : map [string ]string {
792+ "shard1" : ":6379" ,
793+ "shard2" : ":6380" ,
794+ },
795+ })
796+ })
797+
798+ AfterEach (func () {
799+ Expect (ring .Close ()).NotTo (HaveOccurred ())
800+ })
801+
802+ It ("GetShardClients returns active shard clients" , func () {
803+ shards := ring .GetShardClients ()
804+ // Note: This test will pass even if Redis servers are not running,
805+ // because GetShardClients only returns clients that are marked as "up",
806+ // and newly created shards start as "up" until the first health check fails.
807+
808+ if len (shards ) == 0 {
809+ // Expected if Redis servers are not running
810+ Skip ("No active shards found (Redis servers not running)" )
811+ } else {
812+ Expect (len (shards )).To (BeNumerically (">" , 0 ))
813+ for _ , client := range shards {
814+ Expect (client ).NotTo (BeNil ())
815+ }
816+ }
817+ })
818+
819+ It ("GetShardClientForKey returns correct shard for keys" , func () {
820+ testKeys := []string {"key1" , "key2" , "user:123" , "channel:test" }
821+
822+ for _ , key := range testKeys {
823+ client , err := ring .GetShardClientForKey (key )
824+ Expect (err ).NotTo (HaveOccurred ())
825+ Expect (client ).NotTo (BeNil ())
826+ }
827+ })
828+
829+ It ("GetShardClientForKey is consistent for same key" , func () {
830+ key := "test:consistency"
831+
832+ // Call GetShardClientForKey multiple times with the same key
833+ // Should always return the same shard
834+ var firstClient * redis.Client
835+ for i := 0 ; i < 5 ; i ++ {
836+ client , err := ring .GetShardClientForKey (key )
837+ Expect (err ).NotTo (HaveOccurred ())
838+ Expect (client ).NotTo (BeNil ())
839+
840+ if i == 0 {
841+ firstClient = client
842+ } else {
843+ Expect (client .String ()).To (Equal (firstClient .String ()))
844+ }
845+ }
846+ })
847+
848+ It ("GetShardClientForKey distributes keys across shards" , func () {
849+ testKeys := []string {"key1" , "key2" , "key3" , "key4" , "key5" }
850+ shardMap := make (map [string ]int )
851+
852+ for _ , key := range testKeys {
853+ client , err := ring .GetShardClientForKey (key )
854+ Expect (err ).NotTo (HaveOccurred ())
855+ shardMap [client .String ()]++
856+ }
857+
858+ // Should have at least 1 shard (could be all keys go to same shard due to hashing)
859+ Expect (len (shardMap )).To (BeNumerically (">=" , 1 ))
860+ // But with multiple keys, we expect some distribution
861+ Expect (len (shardMap )).To (BeNumerically ("<=" , 2 )) // At most 2 shards (our setup)
862+ })
863+ })
0 commit comments