From f961526cd3686fbbf0f43035f1e7c2afe89350e1 Mon Sep 17 00:00:00 2001 From: Svitlana Kost Date: Wed, 31 Oct 2018 19:16:05 +0200 Subject: [PATCH 1/2] added a count to spop --- aioredis/commands/set.py | 8 +++++--- tests/set_commands_test.py | 30 ++++++++++++++++++++++-------- 2 files changed, 27 insertions(+), 11 deletions(-) diff --git a/aioredis/commands/set.py b/aioredis/commands/set.py index e29ebe5f6..7d153de07 100644 --- a/aioredis/commands/set.py +++ b/aioredis/commands/set.py @@ -43,9 +43,11 @@ def smove(self, sourcekey, destkey, member): """Move a member from one set to another.""" return self.execute(b'SMOVE', sourcekey, destkey, member) - def spop(self, key, *, encoding=_NOTSET): - """Remove and return a random member from a set.""" - return self.execute(b'SPOP', key, encoding=encoding) + def spop(self, key, count=None, *, encoding=_NOTSET): + """Remove and return one or multiple random members from a set.""" + args = [key] + count is not None and args.append(count) + return self.execute(b'SPOP', *args, encoding=encoding) def srandmember(self, key, count=None, *, encoding=_NOTSET): """Get one or multiple random members from a set.""" diff --git a/tests/set_commands_test.py b/tests/set_commands_test.py index d88a12975..b7460d631 100644 --- a/tests/set_commands_test.py +++ b/tests/set_commands_test.py @@ -250,25 +250,39 @@ async def test_smove(redis): @pytest.mark.run_loop async def test_spop(redis): key = b'key:spop:1' - members = b'one', b'two', b'three' - await redis.sadd(key, *members) + members1 = b'one', b'two', b'three' + await redis.sadd(key, *members1) - for _ in members: + for _ in members1: test_result = await redis.spop(key) - assert test_result in members + assert test_result in members1 # test with encoding - members = 'four', 'five', 'six' - await redis.sadd(key, *members) + members2 = 'four', 'five', 'six' + await redis.sadd(key, *members2) - for _ in members: + for _ in members2: test_result = await redis.spop(key, encoding='utf-8') - assert test_result in members + assert test_result in members2 # make sure set is empty, after all values poped test_result = await redis.smembers(key) assert test_result == [] + await redis.sadd(key, *members1) + + # fetch 3 random members + test_result1 = await redis.spop(key, 3) + assert len(test_result1) == 3 + assert set(test_result1).issubset(members1) is True + + await redis.sadd(key, *members2) + + # test with encoding, fetch 3 random members + test_result2 = await redis.spop(key, 3, encoding='utf-8') + assert len(test_result2) == 3 + assert set(test_result2).issubset(members2) is True + # try to pop data from empty set test_result = await redis.spop(b'not:' + key) assert test_result is None From b3d24978f1a14395732e35d829ded46595c0c42d Mon Sep 17 00:00:00 2001 From: Svitlana Kost Date: Fri, 2 Nov 2018 19:26:11 +0200 Subject: [PATCH 2/2] add a count to spop, add tests --- aioredis/commands/set.py | 3 ++- tests/set_commands_test.py | 48 ++++++++++++++++++++++++++++---------- 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/aioredis/commands/set.py b/aioredis/commands/set.py index 7d153de07..6c20b97b1 100644 --- a/aioredis/commands/set.py +++ b/aioredis/commands/set.py @@ -46,7 +46,8 @@ def smove(self, sourcekey, destkey, member): def spop(self, key, count=None, *, encoding=_NOTSET): """Remove and return one or multiple random members from a set.""" args = [key] - count is not None and args.append(count) + if count is not None: + args.append(count) return self.execute(b'SPOP', *args, encoding=encoding) def srandmember(self, key, count=None, *, encoding=_NOTSET): diff --git a/tests/set_commands_test.py b/tests/set_commands_test.py index b7460d631..d3c0ca4f8 100644 --- a/tests/set_commands_test.py +++ b/tests/set_commands_test.py @@ -1,5 +1,7 @@ import pytest +from aioredis import ReplyError + async def add(redis, key, members): ok = await redis.connection.execute(b'sadd', key, members) @@ -250,25 +252,41 @@ async def test_smove(redis): @pytest.mark.run_loop async def test_spop(redis): key = b'key:spop:1' - members1 = b'one', b'two', b'three' - await redis.sadd(key, *members1) + members = b'one', b'two', b'three' + await redis.sadd(key, *members) - for _ in members1: + for _ in members: test_result = await redis.spop(key) - assert test_result in members1 + assert test_result in members # test with encoding - members2 = 'four', 'five', 'six' - await redis.sadd(key, *members2) + members = 'four', 'five', 'six' + await redis.sadd(key, *members) - for _ in members2: + for _ in members: test_result = await redis.spop(key, encoding='utf-8') - assert test_result in members2 + assert test_result in members # make sure set is empty, after all values poped test_result = await redis.smembers(key) assert test_result == [] + # try to pop data from empty set + test_result = await redis.spop(b'not:' + key) + assert test_result is None + + with pytest.raises(TypeError): + await redis.spop(None) + + +@pytest.redis_version( + 3, 2, 0, + reason="The count argument in SPOP is available since redis>=3.2.0" +) +@pytest.mark.run_loop +async def test_spop_count(redis): + key = b'key:spop:1' + members1 = b'one', b'two', b'three' await redis.sadd(key, *members1) # fetch 3 random members @@ -276,6 +294,7 @@ async def test_spop(redis): assert len(test_result1) == 3 assert set(test_result1).issubset(members1) is True + members2 = 'four', 'five', 'six' await redis.sadd(key, *members2) # test with encoding, fetch 3 random members @@ -284,11 +303,16 @@ async def test_spop(redis): assert set(test_result2).issubset(members2) is True # try to pop data from empty set - test_result = await redis.spop(b'not:' + key) - assert test_result is None + test_result = await redis.spop(b'not:' + key, 2) + assert len(test_result) == 0 - with pytest.raises(TypeError): - await redis.spop(None) + # test with negative counter + with pytest.raises(ReplyError): + await redis.spop(key, -2) + + # test with counter is zero + test_result3 = await redis.spop(key, 0) + assert len(test_result3) == 0 @pytest.mark.run_loop