zrange/zrevrange support (#59)

* zrevrange support * fix ZRange and add REV support Co-authored-by: Ivan Ponomarev <iponomarev@mail.ru>
atp-mipt · Mar 13, 2022 · 8072b7e · 8072b7e
1 parent a05a7ff
commit 8072b7e
Show file tree

Hide file tree

Showing 4 changed files with 176 additions and 9 deletions.
diff --git a/src/main/java/com/github/fppt/jedismock/operations/sortedsets/ZRange.java b/src/main/java/com/github/fppt/jedismock/operations/sortedsets/ZRange.java
@@ -1,14 +1,16 @@
 package com.github.fppt.jedismock.operations.sortedsets;
 
 import com.github.fppt.jedismock.datastructures.RMHMap;
+import com.github.fppt.jedismock.datastructures.Slice;
 import com.github.fppt.jedismock.operations.AbstractRedisOperation;
 import com.github.fppt.jedismock.operations.RedisCommand;
 import com.github.fppt.jedismock.server.Response;
-import com.github.fppt.jedismock.datastructures.Slice;
 import com.github.fppt.jedismock.storage.RedisBase;
 
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.ToDoubleFunction;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -18,6 +20,11 @@
 class ZRange extends AbstractRedisOperation {
 
     private static final String WITH_SCORES = "WITHSCORES";
+    private static final String IS_REV = "REV";
+
+    private static final Comparator<Map.Entry<Slice, Double>> zRangeComparator = Comparator
+            .comparingDouble((ToDoubleFunction<Map.Entry<Slice, Double>>) Map.Entry::getValue)
+            .thenComparing(Map.Entry::getKey);
 
     ZRange(RedisBase base, List<Slice> params) {
         super(base, params);
@@ -50,15 +57,28 @@ protected Slice response() {
             end = map.size() - 1;
         }
 
-        final boolean withScores = params().size() == 4 && WITH_SCORES.equalsIgnoreCase(params().get(3).toString());
+        boolean withScores = false;
+        boolean isRev = false;
+
+        for (Slice param : params()) {
+            if (WITH_SCORES.equalsIgnoreCase(param.toString())){
+                withScores = true;
+            }
+            if (IS_REV.equalsIgnoreCase(param.toString())){
+                isRev = true;
+            }
+        }
+
+        boolean finalWithScores = withScores;
         List<Slice> values = map.entrySet().stream()
-            .skip(start)
-            .limit(end - start + 1)
-            .flatMap(e -> withScores
-                    ? Stream.of(e.getKey(), Slice.create(e.getValue().toString()))
-                    : Stream.of(e.getKey()))
-            .map(Response::bulkString)
-            .collect(Collectors.toList());
+                .sorted(isRev ? zRangeComparator.reversed() : zRangeComparator)
+                .skip(start)
+                .limit(end - start + 1)
+                .flatMap(e -> finalWithScores
+                        ? Stream.of(e.getKey(), Slice.create(e.getValue().toString()))
+                        : Stream.of(e.getKey()))
+                .map(Response::bulkString)
+                .collect(Collectors.toList());
 
         return Response.array(values);
     }

diff --git a/src/main/java/com/github/fppt/jedismock/operations/sortedsets/ZRevRange.java b/src/main/java/com/github/fppt/jedismock/operations/sortedsets/ZRevRange.java
@@ -0,0 +1,29 @@
+package com.github.fppt.jedismock.operations.sortedsets;
+
+import com.github.fppt.jedismock.datastructures.Slice;
+import com.github.fppt.jedismock.operations.AbstractRedisOperation;
+import com.github.fppt.jedismock.operations.RedisCommand;
+import com.github.fppt.jedismock.storage.RedisBase;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RedisCommand("zrevrange")
+class ZRevRange extends AbstractRedisOperation {
+
+    private static final String IS_REV = "REV";
+    private final ZRange zRange;
+
+
+    ZRevRange(RedisBase base, List<Slice> params) {
+        super(base, params);
+        List<Slice> updatedParams = new ArrayList<>(params);
+        updatedParams.add(Slice.create(IS_REV));
+        this.zRange = new ZRange(base, updatedParams);
+    }
+
+    @Override
+    protected Slice response() {
+        return zRange.response();
+    }
+}
diff --git a/src/test/java/com/github/fppt/jedismock/comparisontests/sortedsets/TestZRange.java b/src/test/java/com/github/fppt/jedismock/comparisontests/sortedsets/TestZRange.java
@@ -0,0 +1,63 @@
+package com.github.fppt.jedismock.comparisontests.sortedsets;
+
+import com.github.fppt.jedismock.comparisontests.ComparisonBase;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import redis.clients.jedis.Jedis;
+import redis.clients.jedis.params.ZRangeParams;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@ExtendWith(ComparisonBase.class)
+public class TestZRange {
+
+    private static final String ZSET_KEY = "myzset";
+
+    @BeforeEach
+    public void setUp(Jedis jedis) {
+        jedis.flushDB();
+        jedis.zadd(ZSET_KEY, 2, "aaaa");
+        jedis.zadd(ZSET_KEY, 3, "bbbb");
+        jedis.zadd(ZSET_KEY, 1, "cccc");
+        jedis.zadd(ZSET_KEY, 3, "bcbb");
+        jedis.zadd(ZSET_KEY, 3, "babb");
+        assertEquals(5L, jedis.zcount(ZSET_KEY, Integer.MIN_VALUE, Integer.MAX_VALUE));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsEverythingInRightOrderWithPlusMinusMaxInteger(Jedis jedis) {
+        assertEquals(Arrays.asList("cccc", "aaaa", "babb", "bbbb", "bcbb"), new ArrayList<>(jedis.zrange(ZSET_KEY, Integer.MIN_VALUE, Integer.MAX_VALUE)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsListInRightOrderWithPositiveRange(Jedis jedis) {
+        assertEquals(Arrays.asList("aaaa", "babb", "bbbb"), new ArrayList<>(jedis.zrange(ZSET_KEY, 1, 3)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsListInRightOrderWithNegativeRange(Jedis jedis) {
+        assertEquals(Arrays.asList("babb", "bbbb", "bcbb"), new ArrayList<>(jedis.zrange(ZSET_KEY, -3, -1)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsListInRightOrderWithNegativeStartAndPositiveEndRange(Jedis jedis) {
+        assertEquals(Arrays.asList("cccc", "aaaa", "babb"), new ArrayList<>(jedis.zrange(ZSET_KEY, -5, 2)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsListInRightOrderWithPositiveStartAndNegativeEndRange(Jedis jedis) {
+        assertEquals(Arrays.asList("aaaa", "babb", "bbbb", "bcbb"), new ArrayList<>(jedis.zrange(ZSET_KEY, 1, -1)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrange_EnsureItReturnsListInLexicographicOrderForSameScore(Jedis jedis) {
+        jedis.zadd("foo", 42, "def");
+        jedis.zadd("foo", 42, "abc");
+        assertEquals(Arrays.asList("abc", "def"), jedis.zrange("foo", 0, -1));
+        assertEquals(Arrays.asList("def", "abc"), jedis.zrange("foo", ZRangeParams.zrangeParams(0, -1).rev()));
+    }
+}
diff --git a/src/test/java/com/github/fppt/jedismock/comparisontests/sortedsets/TestZRevRange.java b/src/test/java/com/github/fppt/jedismock/comparisontests/sortedsets/TestZRevRange.java
@@ -0,0 +1,55 @@
+package com.github.fppt.jedismock.comparisontests.sortedsets;
+
+import com.github.fppt.jedismock.comparisontests.ComparisonBase;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+import org.junit.jupiter.api.extension.ExtendWith;
+import redis.clients.jedis.Jedis;
+import redis.clients.jedis.exceptions.JedisDataException;
+
+import java.util.*;
+
+import static java.util.Arrays.asList;
+import static org.junit.jupiter.api.Assertions.*;
+
+@ExtendWith(ComparisonBase.class)
+public class TestZRevRange {
+
+    private static final String ZSET_KEY = "myzset";
+
+    @BeforeEach
+    public void setUp(Jedis jedis) {
+        jedis.flushDB();
+        jedis.zadd(ZSET_KEY, 2, "aaaa");
+        jedis.zadd(ZSET_KEY, 3, "bbbb");
+        jedis.zadd(ZSET_KEY, 1, "cccc");
+        jedis.zadd(ZSET_KEY, 3, "bcbb");
+        jedis.zadd(ZSET_KEY, 3, "babb");
+        assertEquals(5L, jedis.zcount(ZSET_KEY, Integer.MIN_VALUE, Integer.MAX_VALUE));
+    }
+
+    @TestTemplate
+    public void whenUsingZrevrange_EnsureItReturnsEverythingInRightOrderWithPlusMinusMaxInteger(Jedis jedis) {
+        assertEquals(Arrays.asList("bcbb", "bbbb", "babb", "aaaa", "cccc"), new ArrayList<>(jedis.zrevrange(ZSET_KEY, Integer.MIN_VALUE, Integer.MAX_VALUE)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrevrange_EnsureItReturnsListInRightOrderWithPositiveRange(Jedis jedis) {
+        assertEquals(Arrays.asList("bbbb", "babb", "aaaa"), new ArrayList<>(jedis.zrevrange(ZSET_KEY, 1, 3)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrevrange_EnsureItReturnsListInRightOrderWithNegativeRange(Jedis jedis) {
+        assertEquals(Arrays.asList("babb", "aaaa", "cccc"), new ArrayList<>(jedis.zrevrange(ZSET_KEY, -3, -1)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrevrange_EnsureItReturnsListInRightOrderWithNegativeStartAndPositiveEndRange(Jedis jedis) {
+        assertEquals(Arrays.asList("bcbb", "bbbb", "babb"), new ArrayList<>(jedis.zrevrange(ZSET_KEY, -5, 2)));
+    }
+
+    @TestTemplate
+    public void whenUsingZrevrange_EnsureItReturnsListInRightOrderWithPositiveStartAndNegativeEndRange(Jedis jedis) {
+        assertEquals(Arrays.asList("bbbb", "babb", "aaaa", "cccc"), new ArrayList<>(jedis.zrevrange(ZSET_KEY, 1, -1)));
+    }
+}