Add test flags to run tests using redis cluster

asynq_test.go

@@ -7,6 +7,7 @@ package asynq
 import (
 	"flag"
 	"sort"
+	"strings"
 	"testing"
 
 	"github.com/go-redis/redis/v7"
@@ -24,6 +25,9 @@ var (
 	redisAddr string
 	redisDB   int
 
+	useRedisCluster   bool
+	redisClusterAddrs string // comma-separated list of host:port
+
 	testLogLevel = FatalLevel
 )
 
@@ -32,23 +36,52 @@ var testLogger *log.Logger
 func init() {
 	flag.StringVar(&redisAddr, "redis_addr", "localhost:6379", "redis address to use in testing")
 	flag.IntVar(&redisDB, "redis_db", 14, "redis db number to use in testing")
+	flag.BoolVar(&useRedisCluster, "redis_cluster", false, "use redis cluster as a broker in testing")
+	flag.StringVar(&redisClusterAddrs, "redis_cluster_addrs", "localhost:7000,localhost:7001,localhost:7002", "comma separated list of redis server addresses")
 	flag.Var(&testLogLevel, "loglevel", "log level to use in testing")
 
 	testLogger = log.NewLogger(nil)
 	testLogger.SetLevel(toInternalLogLevel(testLogLevel))
 }
 
-func setup(tb testing.TB) *redis.Client {
+func setup(tb testing.TB) (r redis.UniversalClient) {
 	tb.Helper()
-	r := redis.NewClient(&redis.Options{
-		Addr: redisAddr,
-		DB:   redisDB,
-	})
+	if useRedisCluster {
+		addrs := strings.Split(redisClusterAddrs, ",")
+		if len(addrs) == 0 {
+			tb.Fatal("No redis cluster addresses provided. Please set addresses using --redis_cluster_addrs flag.")
+		}
+		r = redis.NewClusterClient(&redis.ClusterOptions{
+			Addrs: addrs,
+		})
+	} else {
+		r = redis.NewClient(&redis.Options{
+			Addr: redisAddr,
+			DB:   redisDB,
+		})
+	}
 	// Start each test with a clean slate.
 	h.FlushDB(tb, r)
 	return r
 }
 
+func getRedisConnOpt(tb testing.TB) RedisConnOpt {
+	tb.Helper()
+	if useRedisCluster {
+		addrs := strings.Split(redisClusterAddrs, ",")
+		if len(addrs) == 0 {
+			tb.Fatal("No redis cluster addresses provided. Please set addresses using --redis_cluster_addrs flag.")
+		}
+		return RedisClusterClientOpt{
+			Addrs: addrs,
+		}
+	}
+	return RedisClientOpt{
+		Addr: redisAddr,
+		DB:   redisDB,
+	}
+}
+
 var sortTaskOpt = cmp.Transformer("SortMsg", func(in []*Task) []*Task {
 	out := append([]*Task(nil), in...) // Copy input to avoid mutating it
 	sort.Slice(out, func(i, j int) bool {