Add support for the zrandmember operation and improve srandmember

src/main/java/com/github/fppt/jedismock/Utils.java

@@ -3,6 +3,12 @@
 import com.github.fppt.jedismock.exception.WrongValueTypeException;
 
 import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Random;
 
 /**
  * Created by Xiaolu on 2015/4/21.
@@ -91,6 +97,54 @@ public static String createRegexFromGlob(String glob) {
         return out.toString();
     }
 
+    /**
+     * Performs reservoir sampling on a given collection, returning a random subset of up to {@code count} elements.
+     * <p>
+     * This implementation uses Algorithm L (a fast reservoir sampling method) to efficiently sample elements
+     * with uniform probability without needing to store or shuffle the entire dataset.
+     *
+     * @param collection the input collection to sample from
+     * @param count the number of elements to sample; must be non-negative and less than or equal to the size of the collection
+     * @param random a {@link Random} instance used for generating random numbers
+     * @param <E> the type of elements in the input collection
+     * @return a list containing up to {@code count} randomly sampled elements from the collection
+     * @throws IllegalArgumentException if {@code count} is negative
+     */
+    public static <E> List<E> reservoirSampling(Collection<E> collection, int count, Random random) {
+        if (count < 0) {
+            throw new IllegalArgumentException("count must be non-negative");
+        }
+
+        if (count == 0) {
+            return Collections.emptyList();
+        }
+
+        if (count > collection.size()) {
+            count = collection.size();
+        }
+
+        List<E> result = new ArrayList<>(count);
+        Iterator<E> iter = collection.iterator();
+        for (int i = 0; i < count; i++) {
+            result.add(iter.next());
+        }
+
+        double W = Math.exp(Math.log(random.nextDouble()) / count);
+        while (iter.hasNext()) {
+            int skip = (int)Math.floor(Math.log(random.nextDouble()) / Math.log(1 - W));
+            for (int j = 0; j < skip; j++) {
+                if (!iter.hasNext()) break;
+                iter.next();
+            }
+            if (iter.hasNext()) {
+                result.set(random.nextInt(count), iter.next());
+                W = W * Math.exp(Math.log(random.nextDouble()) / count);
+            }
+        }
+
+        return result;
+    }
+
     public static long toNanoTimeout(String value) {
         return (long) (convertToDouble(value) * 1_000_000_000L);
     }

src/main/java/com/github/fppt/jedismock/operations/sets/SRandMember.java

@@ -1,19 +1,22 @@
 package com.github.fppt.jedismock.operations.sets;
 
+import com.github.fppt.jedismock.Utils;
 import com.github.fppt.jedismock.datastructures.RMSet;
 import com.github.fppt.jedismock.datastructures.Slice;
-import com.github.fppt.jedismock.exception.WrongValueTypeException;
 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.storage.RedisBase;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
+import java.util.Random;
+import java.util.Set;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.stream.Collectors;
 
+import static com.github.fppt.jedismock.Utils.reservoirSampling;
+
 @RedisCommand("srandmember")
 public class SRandMember extends AbstractRedisOperation {
     public SRandMember(RedisBase base, List<Slice> params) {
@@ -28,48 +31,27 @@ protected int minArgs() {
     @Override
     protected Slice response() {
         RMSet set = base().getSet(params().get(0));
-        int number;
-        if (params().size() > 1) {
-            if (set == null) {
-                return Response.EMPTY_ARRAY;
-            }
-            int result;
-            String value = params().get(1).toString();
-            try {
-                result = Integer.parseInt(value);
-            } catch (NumberFormatException e) {
-                throw new WrongValueTypeException("ERR value is out of range");
-            }
-            number = result;
-        } else {
-            if (set == null) {
-                return Response.NULL;
-            }
-            number = 1;
+        boolean isArrayResponse = params().size() > 1;
+
+        int number = isArrayResponse ? Utils.convertToInteger(params().get(1).toString()) : 1;
+        if (set == null) {
+            return isArrayResponse ? Response.EMPTY_ARRAY : Response.NULL;
         }
 
-        // TODO: more effective algorithms should be used here,
-        // avoiding conversion of set to list, shuffling all the elements etc.
-        List<Slice> list = new ArrayList<>(set.getStoredData());
-        if (number == 1) {
-            int index = ThreadLocalRandom.current().nextInt(list.size());
-            return params().size() > 1 ?
-                    Response.array(Response.bulkString(list.get(index))) :
-                    Response.bulkString(list.get(index));
-        } else if (number > 1) {
-            Collections.shuffle(list);
-            return Response.array(
-                    list.stream()
-                            .map(Response::bulkString)
-                            .limit(number)
-                            .collect(Collectors.toList()));
+        List<Slice> result = selectEntries(set.getStoredData(), number, ThreadLocalRandom.current());
+        return isArrayResponse ?
+                Response.array(result.stream().map(Response::bulkString).collect(Collectors.toList())) :
+                Response.bulkString(result.get(0));
+    }
+
+    private List<Slice> selectEntries(Set<Slice> set, int count, Random random) {
+        if (count > 0) {
+            return reservoirSampling(set, count, random);
         } else {
-            List<Slice> result =
-                    ThreadLocalRandom.current().ints(-number, 0, list.size())
-                            .mapToObj(list::get)
-                            .map(Response::bulkString)
-                            .collect(Collectors.toList());
-            return Response.array(result);
+            List<Slice> list = new ArrayList<>(set);
+            return ThreadLocalRandom.current().ints(-count, 0, list.size())
+                    .mapToObj(list::get)
+                    .collect(Collectors.toList());
         }
     }
 }

src/main/java/com/github/fppt/jedismock/operations/sortedsets/ZRandMember.java

@@ -0,0 +1,102 @@
+package com.github.fppt.jedismock.operations.sortedsets;
+
+import com.github.fppt.jedismock.Utils;
+import com.github.fppt.jedismock.datastructures.RMZSet;
+import com.github.fppt.jedismock.datastructures.Slice;
+import com.github.fppt.jedismock.datastructures.ZSetEntry;
+import com.github.fppt.jedismock.exception.WrongValueTypeException;
+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.storage.RedisBase;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static com.github.fppt.jedismock.Utils.reservoirSampling;
+
+
+@RedisCommand("zrandmember")
+public class ZRandMember extends AbstractRedisOperation {
+
+    public ZRandMember(RedisBase base, List<Slice> params) {
+        super(base, params);
+    }
+
+    @Override
+    protected int minArgs() {
+        return 1;
+    }
+
+    @Override
+    protected int maxArgs() {
+        return 3;
+    }
+
+    @Override
+    protected Slice response() {
+        RMZSet set = base().getZSet(params().get(0));
+        boolean isArrayResponse = params().size() > 1;
+
+        int count = parseCount();
+        boolean withScores = parseWithScores();
+        if (set == null) {
+            return isArrayResponse ? Response.EMPTY_ARRAY : Response.NULL;
+        }
+
+        List<ZSetEntry> selectedEntries = selectEntries(set, count);
+        return isArrayResponse
+                ? buildArrayResponse(selectedEntries, withScores)
+                : buildSingleResponse(selectedEntries.get(0));
+    }
+
+    private int parseCount() {
+        if (params().size() >= 2) {
+            return Utils.convertToInteger(params().get(1).toString());
+        }
+        return 1;
+    }
+
+    private boolean parseWithScores() {
+        if (params().size() == 3) {
+            if ("withscores".equalsIgnoreCase(params().get(2).toString())) {
+                return true;
+            } else {
+                throw new WrongValueTypeException("ERR syntax error");
+            }
+        }
+        return false;
+    }
+
+    private List<ZSetEntry> selectEntries(RMZSet set, int count) {
+        if (count > 0) {
+            return reservoirSampling(set.entries(false), count, ThreadLocalRandom.current());
+        } else {
+            List<ZSetEntry> entries = new ArrayList<>(set.entries(false));
+            return ThreadLocalRandom.current()
+                    .ints(-count, 0, entries.size())
+                    .mapToObj(entries::get)
+                    .collect(Collectors.toList());
+        }
+    }
+
+    private Slice buildArrayResponse(List<ZSetEntry> entries, boolean withScores) {
+        return Response.array(entries.stream()
+                .flatMap(e -> withScores
+                        ? Stream.of(e.getValue(), Slice.create(String.valueOf(Math.round(e.getScore()))))
+                        : Stream.of(e.getValue()))
+                .map(Response::bulkString)
+                .collect(Collectors.toList()));
+    }
+
+    private Slice buildSingleResponse(ZSetEntry entry) {
+        return Response.bulkString(entry.getValue());
+    }
+}

src/test/java/com/github/fppt/jedismock/TestUtils.java

@@ -4,6 +4,15 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Random;
+import java.util.Set;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.catchThrowable;
 
 /**
  * Created by Xiaolu on 2015/4/20.
@@ -26,4 +35,77 @@ public void close() throws IOException {
         });
     }
 
+    @Test
+    void testReservoirSamplingBasicFunctionality() {
+        List<Integer> input = Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
+        int count = 5;
+
+        List<Integer> sample = Utils.reservoirSampling(input, count, new Random());
+        Set<Integer> sampleSet = new HashSet<>(sample);
+
+        assertThat(sample).hasSize(count);
+        assertThat(sampleSet).hasSize(count);
+        assertThat(input).containsAll(sample);
+    }
+
+    @Test
+    void testReservoirSamplingWithCountGreaterThanCollection() {
+        List<Integer> input = Arrays.asList(1, 2, 3);
+        int count = 10;
+
+        List<Integer> sample = Utils.reservoirSampling(input, count, new Random());
+        Set<Integer> sampleSet = new HashSet<>(sample);
+
+        assertThat(sample).hasSize(input.size());
+        assertThat(sampleSet).hasSize(input.size());
+        assertThat(input).containsAll(sample);
+    }
+
+    @Test
+    void testReservoirSamplingWithExactCount() {
+        List<Integer> input = Arrays.asList(1, 2, 3);
+        int count = 3;
+
+        List<Integer> sample = Utils.reservoirSampling(input, count, new Random());
+        Set<Integer> sampleSet = new HashSet<>(sample);
+
+        assertThat(sample).hasSize(input.size());
+        assertThat(sampleSet).hasSize(input.size());
+        assertThat(input).containsAll(sample);
+    }
+
+    @Test
+    void testReservoirSamplingWithZeroCount() {
+        List<Integer> input = Arrays.asList(1, 2, 3);
+        int count = 0;
+
+        List<Integer> sample = Utils.reservoirSampling(input, count, new Random());
+
+        assertThat(sample).isEmpty();
+    }
+
+    @Test
+    void testReservoirSamplingWithNegativeCount() {
+        List<Integer> input = Arrays.asList(1, 2, 3);
+        int count = -1;
+
+        Throwable throwable = catchThrowable(() -> Utils.reservoirSampling(input, count, new Random()));
+
+        assertThat(throwable).isInstanceOf(IllegalArgumentException.class);
+    }
+
+    @Test
+    void testReservoirSamplingDeterminismWithFixedSeed() {
+        List<Integer> input = new ArrayList<>();
+        for (int i = 0; i < 100; i++) input.add(i);
+        int count = 10;
+        Random random1 = new Random(123);
+        Random random2 = new Random(123);
+
+        List<Integer> sample1 = Utils.reservoirSampling(input, count, random1);
+        List<Integer> sample2 = Utils.reservoirSampling(input, count, random2);
+
+        assertThat(sample1).isEqualTo(sample2);
+    }
+
 }

src/test/java/com/github/fppt/jedismock/comparisontests/sets/SRandMemberTest.java

@@ -5,6 +5,7 @@
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;
 import redis.clients.jedis.Jedis;
+import redis.clients.jedis.Protocol;
 
 import java.util.Arrays;
 import java.util.Collection;
@@ -102,4 +103,13 @@ void randMemberReturnOneElementAsSingletonList(Jedis jedis) {
         assertThat(jedis.srandmember("key", 0)).isEmpty();
         assertThat(jedis.srandmember("key", -1)).containsExactly("a");
     }
+
+    @TestTemplate
+    void randMemberWithInvalidArgumentAndNonExistingKeyThrowsError(Jedis jedis) {
+        try {
+            jedis.sendCommand(Protocol.Command.SRANDMEMBER, "myzset", "WRONGARG");
+        } catch (Exception e) {
+            assertThat(e.getMessage()).contains("ERR value is not an integer or out of range");
+        }
+    }
 }