Python: Add command SRandMember (valkey-io#1578)

* Python: Add command SRandMember

Co-authored-by: TJ Zhang <[email protected]>

CHANGELOG.md

@@ -37,6 +37,7 @@
 * Node: Added OBJECT IDLETIME command ([#1567](https://github.com/aws/glide-for-redis/pull/1567))
 * Node: Added OBJECT REFCOUNT command ([#1568](https://github.com/aws/glide-for-redis/pull/1568))
 * Python: Added SETBIT command ([#1571](https://github.com/aws/glide-for-redis/pull/1571))
+* Python: Added SRandMember command ([#1578](https://github.com/aws/glide-for-redis/pull/1578))
 * Python: Added GETBIT command ([#1575](https://github.com/aws/glide-for-redis/pull/1575))
 * Python: Added BITCOUNT command ([#1592](https://github.com/aws/glide-for-redis/pull/1592))
 * Python: Added TOUCH command ([#1582](https://github.com/aws/glide-for-redis/pull/1582))

python/python/glide/async_commands/core.py

@@ -4577,3 +4577,55 @@ async def object_refcount(self, key: str) -> Optional[int]:
             Optional[int],
             await self._execute_command(RequestType.ObjectRefCount, [key]),
         )
+
+    async def srandmember(self, key: str) -> Optional[str]:
+        """
+        Returns a random element from the set value stored at 'key'.
+
+        See https://valkey.io/commands/srandmember for more details.
+
+        Args:
+            key (str): The key from which to retrieve the set member.
+
+        Returns:
+            str: A random element from the set, or None if 'key' does not exist.
+
+        Examples:
+            >>> await client.sadd("my_set", {"member1": 1.0, "member2": 2.0})
+            >>> await client.srandmember("my_set")
+            "member1"  # "member1" is a random member of "my_set".
+            >>> await client.srandmember("non_existing_set")
+            None  # "non_existing_set" is not an existing key, so None was returned.
+        """
+        return cast(
+            Optional[str],
+            await self._execute_command(RequestType.SRandMember, [key]),
+        )
+
+    async def srandmember_count(self, key: str, count: int) -> List[str]:
+        """
+        Returns one or more random elements from the set value stored at 'key'.
+
+        See https://valkey.io/commands/srandmember for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            count (int): The number of members to return.
+                If `count` is positive, returns unique members.
+                If `count` is negative, allows for duplicates members.
+
+        Returns:
+            List[str]: A list of members from the set.
+                If the set does not exist or is empty, the response will be an empty list.
+
+        Examples:
+            >>> await client.sadd("my_set", {"member1": 1.0, "member2": 2.0})
+            >>> await client.srandmember("my_set", -3)
+                ["member1", "member1", "member2"]  # "member1" and "member2" are random members of "my_set".
+            >>> await client.srandmember("non_existing_set", 3)
+                []  # "non_existing_set" is not an existing key, so an empty list was returned.
+        """
+        return cast(
+            List[str],
+            await self._execute_command(RequestType.SRandMember, [key, str(count)]),
+        )

python/python/glide/async_commands/transaction.py

@@ -3165,6 +3165,38 @@ def object_refcount(self: TTransaction, key: str) -> TTransaction:
         """
         return self.append_command(RequestType.ObjectRefCount, [key])
 
+    def srandmember(self: TTransaction, key: str) -> TTransaction:
+        """
+        Returns a random element from the set value stored at 'key'.
+
+        See https://valkey.io/commands/srandmember for more details.
+
+        Args:
+            key (str): The key from which to retrieve the set member.
+
+        Command Response:
+            str: A random element from the set, or None if 'key' does not exist.
+        """
+        return self.append_command(RequestType.SRandMember, [key])
+
+    def srandmember_count(self: TTransaction, key: str, count: int) -> TTransaction:
+        """
+        Returns one or more random elements from the set value stored at 'key'.
+
+        See https://valkey.io/commands/srandmember for more details.
+
+        Args:
+            key (str): The key of the sorted set.
+            count (int): The number of members to return.
+                If `count` is positive, returns unique members.
+                If `count` is negative, allows for duplicates members.
+
+        Command Response:
+            List[str]: A list of members from the set.
+                If the set does not exist or is empty, the response will be an empty list.
+        """
+        return self.append_command(RequestType.SRandMember, [key, str(count)])
+
 
 class Transaction(BaseTransaction):
     """

python/python/tests/test_async_client.py

@@ -5013,6 +5013,51 @@ async def test_object_refcount(self, redis_client: TRedisClient):
         refcount = await redis_client.object_refcount(string_key)
         assert refcount is not None and refcount >= 0
 
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_srandmember(self, redis_client: TRedisClient):
+        key = get_random_string(10)
+        string_key = get_random_string(10)
+        elements = ["one", "two"]
+        assert await redis_client.sadd(key, elements) == 2
+
+        member = await redis_client.srandmember(key)
+        assert member in elements
+        assert await redis_client.srandmember("non_existing_key") is None
+
+        # key exists, but it is not a set
+        assert await redis_client.set(string_key, "value") == OK
+        with pytest.raises(RequestError):
+            await redis_client.srandmember(string_key)
+
+    @pytest.mark.parametrize("cluster_mode", [True, False])
+    @pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
+    async def test_srandmember_count(self, redis_client: TRedisClient):
+        key = get_random_string(10)
+        string_key = get_random_string(10)
+        elements = ["one", "two"]
+        assert await redis_client.sadd(key, elements) == 2
+
+        # unique values are expected as count is positive
+        members = await redis_client.srandmember_count(key, 4)
+        assert len(members) == 2
+        assert set(members) == {"one", "two"}
+
+        # duplicate values are expected as count is negative
+        members = await redis_client.srandmember_count(key, -4)
+        assert len(members) == 4
+        for member in members:
+            assert member in elements
+
+        # empty return values for non-existing or empty keys
+        assert await redis_client.srandmember_count(key, 0) == []
+        assert await redis_client.srandmember_count("non_existing_key", 0) == []
+
+        # key exists, but it is not a set
+        assert await redis_client.set(string_key, "value") == OK
+        with pytest.raises(RequestError):
+            await redis_client.srandmember_count(string_key, 8)
+
 
 class TestMultiKeyCommandCrossSlot:
     @pytest.mark.parametrize("cluster_mode", [True])

python/python/tests/test_transaction.py

@@ -448,6 +448,12 @@ async def transaction_test(
         alpha=True,
     )
     args.append(4)
+    transaction.sadd(key7, ["one"])
+    args.append(1)
+    transaction.srandmember(key7)
+    args.append("one")
+    transaction.srandmember_count(key7, 1)
+    args.append(["one"])
 
     min_version = "7.0.0"
     if not await check_if_server_version_lt(redis_client, min_version):