Repository: fppt/jedis-mock
Pull request #41 (2019-05-04 22:20:55)

Adding Basic Sorted Set Operations (zadd, zrem, zrange)
The Sorted Set is stored as a LinkedHashMap, with the map's key being the value in the set, and the map's value being the score. This ensures that the values inserted into the set are unique (map's key is unique), and we can easily update the score for a given value. Caveat: every time we insert a new value (or update a value's score), we have to re-sort the LinkedHashMap.


Patch:
diff --git a/src/main/java/com/github/fppt/jedismock/Utils.java b/src/main/java/com/github/fppt/jedismock/Utils.java
index 82c03f8..7aa9a7f 100644
--- a/src/main/java/com/github/fppt/jedismock/Utils.java
+++ b/src/main/java/com/github/fppt/jedismock/Utils.java
@@ -88,6 +88,14 @@ public class Utils {
         }
     }
 
+    public static double convertToDouble(String value){
+        try {
+            return Double.parseDouble(value);
+        } catch (NumberFormatException e) {
+            throw new WrongValueTypeException("ERR bit offset is not a double or out of range");
+        }
+    }
+
     public static String createRegexFromGlob(String glob)
     {
         StringBuilder out = new StringBuilder("^");
diff --git a/src/main/java/com/github/fppt/jedismock/operations/OperationFactory.java b/src/main/java/com/github/fppt/jedismock/operations/OperationFactory.java
index b575f9c..4d40f14 100644
--- a/src/main/java/com/github/fppt/jedismock/operations/OperationFactory.java
+++ b/src/main/java/com/github/fppt/jedismock/operations/OperationFactory.java
@@ -71,6 +71,9 @@ public class OperationFactory {
         TRANSACTIONAL_OPERATIONS.put("smembers", RO_smembers::new);
         TRANSACTIONAL_OPERATIONS.put("hsetnx", RO_hsetnx::new);
         TRANSACTIONAL_OPERATIONS.put("time", RO_time::new);
+        TRANSACTIONAL_OPERATIONS.put("zadd", RO_zadd::new);
+        TRANSACTIONAL_OPERATIONS.put("zrange", RO_zrange::new);
+        TRANSACTIONAL_OPERATIONS.put("zrem", RO_zrem::new);
     }
 
 
diff --git a/src/main/java/com/github/fppt/jedismock/operations/RO_zadd.java b/src/main/java/com/github/fppt/jedismock/operations/RO_zadd.java
new file mode 100644
index 0000000..916df9a
--- /dev/null
+++ b/src/main/java/com/github/fppt/jedismock/operations/RO_zadd.java
@@ -0,0 +1,55 @@
+package com.github.fppt.jedismock.operations;
+
+import com.github.fppt.jedismock.server.Response;
+import com.github.fppt.jedismock.server.Slice;
+import com.github.fppt.jedismock.storage.RedisBase;
+
+import java.util.AbstractMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+import static com.github.fppt.jedismock.Utils.convertToDouble;
+import static com.github.fppt.jedismock.Utils.serializeObject;
+import static java.util.stream.Collectors.toMap;
+
+class RO_zadd extends AbstractRedisOperation {
+
+    RO_zadd(RedisBase base, List<Slice> params) {
+        super(base, params);
+    }
+
+    @Override
+    Slice response() {
+        Slice key = params().get(0);
+        LinkedHashMap<Slice, Slice> map = getDataFromBase(key, new LinkedHashMap<>());
+
+        int count = 0;
+        for (int i = 1; i < params().size(); i += 2) {
+            Slice score = params().get(i);
+            Slice value = params().get(i + 1);
+
+            // Score must be a double. Will throw an exception if it's not.
+            convertToDouble(score.toString());
+
+            Slice prevScore = map.put(value, score);
+            if (prevScore == null) count++;
+        }
+
+        // Sort the map by value (the score) before saving
+        Map<Slice, Double> sortedMap = map.entrySet()
+            .stream()
+            .map(x -> new AbstractMap.SimpleEntry<>(x.getKey(), convertToDouble(x.getValue().toString())))
+            .sorted(Map.Entry.comparingByValue())
+            .collect(
+                toMap(Map.Entry::getKey, Map.Entry::getValue, (e1, e2) -> e2, LinkedHashMap::new));
+
+        try {
+            base().putValue(key, serializeObject(sortedMap));
+        } catch (Exception e) {
+            throw new RuntimeException(e.getMessage());
+        }
+        return Response.integer(count);
+    }
+
+}
diff --git a/src/main/java/com/github/fppt/jedismock/operations/RO_zrange.java b/src/main/java/com/github/fppt/jedismock/operations/RO_zrange.java
new file mode 100644
index 0000000..8f7b0bb
--- /dev/null
+++ b/src/main/java/com/github/fppt/jedismock/operations/RO_zrange.java
@@ -0,0 +1,54 @@
+package com.github.fppt.jedismock.operations;
+
+import com.github.fppt.jedismock.server.Response;
+import com.github.fppt.jedismock.server.Slice;
+import com.github.fppt.jedismock.storage.RedisBase;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import static com.github.fppt.jedismock.Utils.convertToInteger;
+
+class RO_zrange extends AbstractRedisOperation {
+
+    RO_zrange(RedisBase base, List<Slice> params) {
+        super(base, params);
+    }
+
+    @Override
+    Slice response() {
+        Slice key = params().get(0);
+        Map<Slice, Slice> map = getDataFromBase(key, new LinkedHashMap<>());
+
+        int start = convertToInteger(params().get(1).toString());
+        int end = convertToInteger(params().get(2).toString());
+
+        if (start < 0) {
+            start = map.size() + start;
+            if (start < 0) {
+                start = 0;
+            }
+        }
+
+        if (end < 0) {
+            end = map.size() + end;
+            if (end < 0) {
+                end = -1;
+            }
+        }
+
+        if (end >= map.size()) {
+            end = map.size() - 1;
+        }
+
+        List<Slice> values = map.keySet().stream()
+            .skip(start)
+            .limit(end - start + 1)
+            .map(Response::bulkString)
+            .collect(Collectors.toList());
+
+        return Response.array(values);
+    }
+}
diff --git a/src/main/java/com/github/fppt/jedismock/operations/RO_zrem.java b/src/main/java/com/github/fppt/jedismock/operations/RO_zrem.java
new file mode 100644
index 0000000..d995168
--- /dev/null
+++ b/src/main/java/com/github/fppt/jedismock/operations/RO_zrem.java
@@ -0,0 +1,35 @@
+package com.github.fppt.jedismock.operations;
+
+import com.github.fppt.jedismock.server.Response;
+import com.github.fppt.jedismock.server.Slice;
+import com.github.fppt.jedismock.storage.RedisBase;
+
+import java.util.List;
+import java.util.Map;
+
+import static com.github.fppt.jedismock.Utils.serializeObject;
+
+class RO_zrem extends AbstractRedisOperation {
+
+    RO_zrem(RedisBase base, List<Slice> params) {
+        super(base, params);
+    }
+
+    Slice response() {
+        Slice key = params().get(0);
+        Map<Slice, Slice> map = getDataFromBase(key, null);
+        if(map == null || map.isEmpty()) return Response.integer(0);
+        int count = 0;
+        for (int i = 1; i < params().size(); i++) {
+            if (map.remove(params().get(i)) != null) {
+                count++;
+            }
+        }
+        try {
+            base().putValue(key, serializeObject(map));
+        } catch (Exception e) {
+            throw new RuntimeException(e.getMessage());
+        }
+        return Response.integer(count);
+    }
+}


Test patch:
diff --git a/src/test/java/com/github/fppt/jedismock/comparisontests/SimpleOperationsTest.java b/src/test/java/com/github/fppt/jedismock/comparisontests/SimpleOperationsTest.java
index 63b6e6b..f3babc4 100644
--- a/src/test/java/com/github/fppt/jedismock/comparisontests/SimpleOperationsTest.java
+++ b/src/test/java/com/github/fppt/jedismock/comparisontests/SimpleOperationsTest.java
@@ -491,4 +491,120 @@ public class SimpleOperationsTest extends ComparisonBase {
 
         assertNotEquals(ScanParams.SCAN_POINTER_START, result.getStringCursor());
     }
+
+    @Theory
+    public void zaddAddsKey(Jedis jedis) {
+
+        jedis.flushDB();
+
+        String key = "mykey";
+        double score = 10;
+        String value = "myvalue";
+
+        long result = jedis.zadd(key, score, value);
+
+        assertEquals(1L, result);
+
+        List<String> results = new LinkedList<>(jedis.zrange(key, 0, -1));
+
+        assertEquals(1, results.size());
+        assertEquals(value, results.get(0));
+    }
+
+    @Theory
+    public void zaddAddsKeys(Jedis jedis) {
+
+        jedis.flushDB();
+
+        String key = "mykey";
+        Map<String, Double> members = new HashMap<>();
+        members.put("myvalue1", 10d);
+        members.put("myvalue2", 20d);
+
+        long result = jedis.zadd(key, members);
+
+        assertEquals(2L, result);
+
+        List<String> results = new LinkedList<>(jedis.zrange(key, 0, -1));
+
+        assertEquals(2, results.size());
+        assertEquals("myvalue1", results.get(0));
+        assertEquals("myvalue2", results.get(1));
+    }
+
+    @Theory
+    public void zremRemovesKey(Jedis jedis) {
+
+        jedis.flushDB();
+
+        String key = "mykey";
+        Map<String, Double> members = new HashMap<>();
+        members.put("myvalue1", 10d);
+        members.put("myvalue2", 20d);
+
+        long result = jedis.zadd(key, members);
+
+        assertEquals(2L, result);
+
+        List<String> results = new LinkedList<>(jedis.zrange(key, 0, -1));
+
+        assertEquals(2, results.size());
+        assertEquals("myvalue1", results.get(0));
+        assertEquals("myvalue2", results.get(1));
+
+        result = jedis.zrem(key, "myvalue1");
+
+        assertEquals(1L, result);
+
+        results = new LinkedList<>(jedis.zrange(key, 0, -1));
+
+        assertEquals(1, results.size());
+        assertEquals("myvalue2", results.get(0));
+    }
+
+    @Theory
+    public void zrangeKeysCorrectOrder(Jedis jedis) {
+
+        jedis.flushDB();
+
+        String key = "mykey";
+        Map<String, Double> members = new HashMap<>();
+        members.put("myvalue2", 10d);
+        members.put("myvalue4", 20d);
+        members.put("myvalue3", 15d);
+        members.put("myvalue1", 9d);
+
+        long result = jedis.zadd(key, members);
+
+        assertEquals(4L, result);
+
+        List<String> results = new LinkedList<>(jedis.zrange(key, 0, -1));
+
+        assertEquals(4, results.size());
+        assertEquals("myvalue1", results.get(0));
+        assertEquals("myvalue2", results.get(1));
+        assertEquals("myvalue3", results.get(2));
+        assertEquals("myvalue4", results.get(3));
+    }
+
+    @Theory
+    public void zrangeIndexOutOfRange(Jedis jedis) {
+
+        jedis.flushDB();
+
+        String key = "mykey";
+        Map<String, Double> members = new HashMap<>();
+        members.put("myvalue2", 10d);
+        members.put("myvalue4", 20d);
+        members.put("myvalue3", 15d);
+        members.put("myvalue1", 9d);
+
+        long result = jedis.zadd(key, members);
+
+        assertEquals(4L, result);
+
+        Set<String> results = jedis.zrange(key, 0, -6);
+
+        assertEquals(0, results.size());
+    }
 }