Add support for GET option to SET

bmerry · bmerry · commit 93d4ed6266e4 · 2021-11-15T11:31:19.000+02:00
I needed to implement this after switching to validating against a 6.2
Redis server, because the fuzzer in the hypothesis tests kept testing
that feature by accident.
diff --git a/fakeredis/_server.py b/fakeredis/_server.py
@@ -1497,6 +1497,7 @@ def set_(self, key, value, *args):
         xx = False
         nx = False
         keepttl = False
+        get = False
         while i < len(args):
             if casematch(args[i], b'nx'):
                 nx = True
@@ -1517,15 +1518,27 @@ def set_(self, key, value, *args):
             elif casematch(args[i], b'keepttl'):
                 keepttl = True
                 i += 1
+            elif casematch(args[i], b'get'):
+                get = True
+                i += 1
             else:
                 raise SimpleError(SYNTAX_ERROR_MSG)
         if (xx and nx) or ((px is not None) + (ex is not None) + keepttl > 1):
             raise SimpleError(SYNTAX_ERROR_MSG)
+        if nx and get:
+            # The command docs say this is allowed from Redis 7.0.
+            raise SimpleError(SYNTAX_ERROR_MSG)
+
+        old_value = None
+        if get:
+            if key.value is not None and type(key.value) is not bytes:
+                raise SimpleError(WRONGTYPE_MSG)
+            old_value = key.value
 
         if nx and key:
-            return None
+            return old_value
         if xx and not key:
-            return None
+            return old_value
         if not keepttl:
             key.value = value
         else:
@@ -1534,7 +1547,7 @@ def set_(self, key, value, *args):
             key.expireat = self._db.time + ex
         if px is not None:
             key.expireat = self._db.time + px / 1000.0
-        return OK
+        return OK if not get else old_value
 
     @command((Key(), Int, bytes))
     def setex(self, key, seconds, value):
diff --git a/test/test_fakeredis.py b/test/test_fakeredis.py
@@ -821,6 +821,38 @@ def test_set_xx(r):
     assert r.set('foo', 'bar', xx=True) is True
 
 
+@pytest.mark.min_server('6.2')
+def test_set_get(r):
+    assert raw_command(r, 'set', 'foo', 'bar', 'GET') is None
+    assert r.get('foo') == b'bar'
+    assert raw_command(r, 'set', 'foo', 'baz', 'GET') == b'bar'
+    assert r.get('foo') == b'baz'
+
+
+@pytest.mark.min_server('6.2')
+def test_set_get_xx(r):
+    assert raw_command(r, 'set', 'foo', 'bar', 'XX', 'GET') is None
+    assert r.get('foo') is None
+    r.set('foo', 'bar')
+    assert raw_command(r, 'set', 'foo', 'baz', 'XX', 'GET') == b'bar'
+    assert r.get('foo') == b'baz'
+    assert raw_command(r, 'set', 'foo', 'baz', 'GET') == b'baz'
+
+
+@pytest.mark.min_server('6.2')
+def test_set_get_nx(r):
+    # Note: this will most likely fail on a 7.0 server, based on the docs for SET
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'set', 'foo', 'bar', 'NX', 'GET')
+
+
+@pytest.mark.min_server('6.2')
+def set_get_wrongtype(r):
+    r.lpush('foo', 'bar')
+    with pytest.raises(redis.ResponseError):
+        raw_command(r, 'set', 'foo', 'bar', 'GET')
+
+
 def test_del_operator(r):
     r['foo'] = 'bar'
     del r['foo']
diff --git a/test/test_hypothesis.py b/test/test_hypothesis.py
@@ -204,7 +204,10 @@ def commands(*args, **kwargs):
     | commands(st.just('mget'), st.lists(keys))
     | commands(st.sampled_from(['mset', 'msetnx']), st.lists(st.tuples(keys, values)))
     | commands(st.just('set'), keys, values,
-               st.none() | st.just('nx'), st.none() | st.just('xx'))
+               st.none() | st.just('nx'),
+               st.none() | st.just('xx'),
+               st.none() | st.just('keepttl'),
+               st.none() | st.just('get'))
     | commands(st.just('setex'), keys, expires_seconds, values)
     | commands(st.just('psetex'), keys, expires_ms, values)
     | commands(st.just('setnx'), keys, values)
