Implement BZPOPMIN/BZPOPMAX support (#60)

REDIS_COMMANDS.md

@@ -60,6 +60,33 @@ list of [unimplemented commands](#unimplemented-commands).
  * rpush
  * rpushx
 
+### sorted-set
+ * bzpopmax
+ * bzpopmin
+ * zadd
+ * zcard
+ * zcount
+ * zincrby
+ * zinterstore
+ * zlexcount
+ * zpopmax
+ * zpopmin
+ * zrange
+ * zrangebylex
+ * zrangebyscore
+ * zrank
+ * zrem
+ * zremrangebylex
+ * zremrangebyrank
+ * zremrangebyscore
+ * zrevrange
+ * zrevrangebylex
+ * zrevrangebyscore
+ * zrevrank
+ * zscan
+ * zscore
+ * zunionstore
+
 ### generic
  * del
  * dump
@@ -148,33 +175,8 @@ list of [unimplemented commands](#unimplemented-commands).
  * sunion
  * sunionstore
 
-### sorted-set
- * zadd
- * zcard
- * zcount
- * zincrby
- * zinterstore
- * zlexcount
- * zpopmax
- * zpopmin
- * zrange
- * zrangebylex
- * zrangebyscore
- * zrank
- * zrem
- * zremrangebylex
- * zremrangebyrank
- * zremrangebyscore
- * zrevrange
- * zrevrangebylex
- * zrevrangebyscore
- * zrevrank
- * zscan
- * zscore
- * zunionstore
-
 # Unimplemented Commands
-All of the redis commands are implemented in fakeredis with these exceptions:
+All the redis commands are implemented in fakeredis with these exceptions:
 
 ### server
  * acl
@@ -316,8 +318,6 @@ All of the redis commands are implemented in fakeredis with these exceptions:
 
 ### sorted-set
  * bzmpop
- * bzpopmax
- * bzpopmin
  * zdiff
  * zdiffstore
  * zinter

fakeredis/_basefakesocket.py

@@ -6,8 +6,7 @@
 import redis
 
 from . import _msgs as msgs
-from ._commands import (
-    Int)
+from ._commands import (Int, CommandItem)
 from ._helpers import (
     SimpleError, valid_response_type, SimpleString, NoResponse, casematch,
     compile_pattern, QUEUED)
@@ -473,3 +472,11 @@ def _zpop(self, key, count, reverse):
         for item in members:
             zset.discard(item)
         return res
+
+    def _bzpop(self, keys, reverse, first_pass):
+        for key in keys:
+            item = CommandItem(key, self._db, item=self._db.get(key), default=[])
+            temp_res = self._zpop(item, 1, reverse)
+            if temp_res:
+                return [key, temp_res[0], temp_res[1]]
+        return None

fakeredis/_basefakeluasupport.py → fakeredis/_fakeluasocket.py

@@ -35,7 +35,7 @@ def _lua_redis_log(lua_runtime, expected_globals, lvl, *args):
     LOGGER.log(REDIS_LOG_LEVELS_TO_LOGGING[lvl], msg)
 
 
-class BaseFakeLuaSocket:
+class FakeLuaSocket:
 
     # Script commands
     # script debug and script kill will probably not be supported

fakeredis/_fakesocket.py

@@ -9,18 +9,18 @@
 import redis
 
 from . import _msgs as msgs
-from ._basefakeluasupport import BaseFakeLuaSocket
 from ._basefakesocket import BaseFakeSocket
 from ._commands import (
     Key, command, DbIndex, Int, CommandItem, BeforeAny, SortFloat, Float, BitOffset, BitValue, Hash,
     StringTest, ScoreTest, Timeout)
+from ._fakeluasocket import FakeLuaSocket
 from ._helpers import (
     PONG, OK, MAX_STRING_SIZE, SimpleError, SimpleString, casematch,
     BGSAVE_STARTED, casenorm, compile_pattern)
 from ._zset import ZSet
 
 
-class FakeSocket(BaseFakeSocket, BaseFakeLuaSocket):
+class FakeSocket(BaseFakeSocket, FakeLuaSocket):
     _connection_error_class = redis.ConnectionError
 
     def __init__(self, server):
@@ -1093,7 +1093,17 @@ def zpopmin(self, key, count=1):
     def zpopmax(self, key, count=1):
         return self._zpop(key, count, True)
 
-    # TODO: bzpopmin/bzpopmax,
+    @command((bytes, bytes), (bytes,), flags='s')
+    def bzpopmin(self, *args):
+        keys = args[:-1]
+        timeout = Timeout.decode(args[-1])
+        return self._blocking(timeout, functools.partial(self._bzpop, keys, False))
+
+    @command((bytes, bytes), (bytes,), flags='s')
+    def bzpopmax(self, *args):
+        keys = args[:-1]
+        timeout = Timeout.decode(args[-1])
+        return self._blocking(timeout, functools.partial(self._bzpop, keys, True))
 
     @staticmethod
     def _limit_items(items, offset, count):

test/test_sorted_set_cmds.py

@@ -4,6 +4,9 @@
 import testtools
 
 REDIS_VERSION = Version(redis.__version__)
+pytestmark = [
+    testtools.run_test_if_redispy_ver('above', '3'),
+]
 
 
 def test_zpopmin(r):
@@ -35,3 +38,20 @@ def test_zpopmax_too_many(r):
     testtools.zadd(r, 'foo', {'three': 3})
     assert r.zpopmax('foo', count=5) == [(b'three', 3.0), (b'two', 2.0), (b'one', 1.0), ]
 
+
+def test_bzpopmin(r):
+    testtools.zadd(r, 'foo', {'one': 1, 'two': 2, 'three': 3})
+    testtools.zadd(r, 'bar', {'a': 1.5, 'b': 2, 'c': 3})
+    assert r.bzpopmin(['foo', 'bar'], 0) == (b'foo', b'one', 1.0)
+    assert r.bzpopmin(['foo', 'bar'], 0) == (b'foo', b'two', 2.0)
+    assert r.bzpopmin(['foo', 'bar'], 0) == (b'foo', b'three', 3.0)
+    assert r.bzpopmin(['foo', 'bar'], 0) == (b'bar', b'a', 1.5)
+
+
+def test_bzpopmax(r):
+    testtools.zadd(r, 'foo', {'one': 1, 'two': 2, 'three': 3})
+    testtools.zadd(r, 'bar', {'a': 1.5, 'b': 2.5, 'c': 3.5})
+    assert r.bzpopmax(['foo', 'bar'], 0) == (b'foo', b'three', 3.0)
+    assert r.bzpopmax(['foo', 'bar'], 0) == (b'foo', b'two', 2.0)
+    assert r.bzpopmax(['foo', 'bar'], 0) == (b'foo', b'one', 1.0)
+    assert r.bzpopmax(['foo', 'bar'], 0) == (b'bar', b'c', 3.5)