Python: add SUNIONSTORE command (#286)

CHANGELOG.md

@@ -28,6 +28,7 @@
 * Python: Added BZPOPMIN and BZPOPMAX commands ([#1399](https://github.com/aws/glide-for-redis/pull/1399))
 * Python: Added ZUNIONSTORE, ZINTERSTORE commands ([#1388](https://github.com/aws/glide-for-redis/pull/1388))
 * Python: Added ZRANDMEMBER command ([#1413](https://github.com/aws/glide-for-redis/pull/1413))
+* Python: Added SUNIONSTORE command ([#1423](https://github.com/aws/glide-for-redis/pull/1423))
 
 
 #### Fixes

python/python/glide/async_commands/core.py

@@ -1527,6 +1527,36 @@ async def sismember(
             await self._execute_command(RequestType.SIsMember, [key, member]),
         )
 
+    async def sunionstore(
+        self,
+        destination: str,
+        keys: List[str],
+    ) -> int:
+        """
+        Stores the members of the union of all given sets specified by `keys` into a new set at `destination`.
+
+        When in cluster mode, all keys in `keys` and `destination` must map to the same hash slot.
+
+        See https://valkey.io/commands/sunionstore for more details.
+
+        Args:
+            destination (str): The key of the destination set.
+            keys (List[str]): The keys from which to retrieve the set members.
+
+        Returns:
+            int: The number of elements in the resulting set.
+
+        Examples:
+            >>> await client.sadd("set1", ["member1"])
+            >>> await client.sadd("set2", ["member2"])
+            >>> await client.sunionstore("my_set", ["set1", "set2"])
+                2  # Two elements were stored in "my_set", and those two members are the union of "set1" and "set2".
+        """
+        return cast(
+            int,
+            await self._execute_command(RequestType.SUnionStore, [destination] + keys),
+        )
+
     async def ltrim(self, key: str, start: int, end: int) -> TOK:
         """
         Trim an existing list so that it will contain only the specified range of elements specified.

python/python/glide/async_commands/transaction.py

@@ -1010,6 +1010,25 @@ def sismember(
         """
         return self.append_command(RequestType.SIsMember, [key, member])
 
+    def sunionstore(
+        self: TTransaction,
+        destination: str,
+        keys: List[str],
+    ) -> TTransaction:
+        """
+        Stores the members of the union of all given sets specified by `keys` into a new set at `destination`.
+
+        See https://valkey.io/commands/sunionstore for more details.
+
+        Args:
+            destination (str): The key of the destination set.
+            keys (List[str]): The keys from which to retrieve the set members.
+
+        Command response:
+            int: The number of elements in the resulting set.
+        """
+        return self.append_command(RequestType.SUnionStore, [destination] + keys)
+
     def ltrim(self: TTransaction, key: str, start: int, end: int) -> TTransaction:
         """
         Trim an existing list so that it will contain only the specified range of elements specified.

python/python/tests/test_async_client.py

@@ -1107,6 +1107,59 @@ async def test_spop(self, redis_client: TRedisClient):
         assert await redis_client.spop("non_existing_key") == None
         assert await redis_client.spop_count("non_existing_key", 3) == set()
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_sunionstore(self, redis_client: TRedisClient):
+        key1 = f"{{testKey}}:1-{get_random_string(10)}"
+        key2 = f"{{testKey}}:2-{get_random_string(10)}"
+        key3 = f"{{testKey}}:3-{get_random_string(10)}"
+        key4 = f"{{testKey}}:4-{get_random_string(10)}"
+        string_key = f"{{testKey}}:4-{get_random_string(10)}"
+        non_existing_key = f"{{testKey}}:5-{get_random_string(10)}"
+
+        assert await redis_client.sadd(key1, ["a", "b", "c"]) == 3
+        assert await redis_client.sadd(key2, ["c", "d", "e"]) == 3
+        assert await redis_client.sadd(key3, ["e", "f", "g"]) == 3
+
+        # store union in new key
+        assert await redis_client.sunionstore(key4, [key1, key2]) == 5
+        assert await redis_client.smembers(key4) == {"a", "b", "c", "d", "e"}
+
+        # overwrite existing set
+        assert await redis_client.sunionstore(key1, [key4, key2]) == 5
+        assert await redis_client.smembers(key1) == {"a", "b", "c", "d", "e"}
+
+        # overwrite one of the source keys
+        assert await redis_client.sunionstore(key2, [key4, key2]) == 5
+        assert await redis_client.smembers(key1) == {"a", "b", "c", "d", "e"}
+
+        # union with a non existing key
+        assert await redis_client.sunionstore(key2, [non_existing_key]) == 0
+        assert await redis_client.smembers(key2) == set()
+
+        # key exists, but it is not a sorted set
+        assert await redis_client.set(string_key, "value") == OK
+        with pytest.raises(RequestError):
+            await redis_client.sunionstore(key4, [string_key, key1])
+
+        # overwrite destination when destination is not a set
+        assert await redis_client.sunionstore(string_key, [key1, key3]) == 7
+        assert await redis_client.smembers(string_key) == {
+            "a",
+            "b",
+            "c",
+            "d",
+            "e",
+            "f",
+            "g",
+        }
+
+        # same-slot requirement
+        if isinstance(redis_client, RedisClusterClient):
+            with pytest.raises(RequestError) as e:
+                await redis_client.sunionstore("abc", ["zxy", "lkn"])
+            assert "CrossSlot" in str(e)
+
     @pytest.mark.parametrize("cluster_mode", [True, False])
     @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
     async def test_ltrim(self, redis_client: TRedisClient):

python/python/tests/test_transaction.py

@@ -201,6 +201,8 @@ async def transaction_test(
     args.append("bar")
     transaction.sadd(key7, ["foo", "bar"])
     args.append(2)
+    transaction.sunionstore(key7, [key7, key7])
+    args.append(2)
     transaction.spop_count(key7, 4)
     args.append({"foo", "bar"})