From 4b228cf7316e3370126df149187c4148e0177089 Mon Sep 17 00:00:00 2001
From: "zach.lee" <zach.lee@sendbird.com>
Date: Fri, 30 Jun 2023 22:23:18 +0900
Subject: [PATCH] make ZUNION, ZUNIONSTORE to avoid command_keys as well

---
 redis/cluster.py      |  8 ++++++++
 tests/test_cluster.py | 20 ++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/redis/cluster.py b/redis/cluster.py
index f19408ad8d..20b353931e 100644
--- a/redis/cluster.py
+++ b/redis/cluster.py
@@ -968,6 +968,14 @@ def determine_slot(self, *args):
             if len(eval_keys) == 0:
                 return random.randrange(0, REDIS_CLUSTER_HASH_SLOTS)
             keys = eval_keys
+        elif command == "ZUNIONSTORE":
+            # https://github.com/redis/redis/pull/12380
+            # command syntax: ZUNIONSTORE destination numkeys key [key ...]
+            #   [WEIGHTS weight [weight ...]] [AGGREGATE <SUM | MIN | MAX>]
+            if len(args) <= 3:
+                raise RedisClusterException(f"Invalid args in ZUNIONSTORE: {args}")
+            num_actual_keys = args[2]
+            keys = (args[1],) + args[3 : 3 + num_actual_keys]
         else:
             keys = self._get_command_keys(*args)
             if keys is None or len(keys) == 0:
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index e3b43fc4be..890b844cea 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -2006,10 +2006,30 @@ def test_cluster_zunion(self, r):
         ]
 
     def test_cluster_zunionstore_sum(self, r):
+        assert r.zunionstore("{foo}d", ["{foo}" + str(i) for i in range(0, 256)]) == 0
+
         r.zadd("{foo}a", {"a1": 1, "a2": 1, "a3": 1})
         r.zadd("{foo}b", {"a1": 2, "a2": 2, "a3": 2})
         r.zadd("{foo}c", {"a1": 6, "a3": 5, "a4": 4})
+
+        result_key = "{foo}d"
+        failed_keys = ["{foo1}a", "{foo}b", "{foo}c"]
+        with pytest.raises(
+            RedisClusterException,
+            match="ZUNIONSTORE - all keys must map to the same key slot",
+        ):
+            r.zunionstore(result_key, failed_keys)
+
+        result_key = "{foo1}d"
+        failed_keys = ["{foo}a", "{foo}b"]
+        with pytest.raises(
+            RedisClusterException,
+            match="ZUNIONSTORE - all keys must map to the same key slot",
+        ):
+            r.zunionstore(result_key, failed_keys)
+
         assert r.zunionstore("{foo}d", ["{foo}a", "{foo}b", "{foo}c"]) == 4
+        assert r.zunionstore("{foo}e", ["{foo}a"]) == 3
         assert r.zrange("{foo}d", 0, -1, withscores=True) == [
             (b"a2", 3),
             (b"a4", 4),