allow srandmembers to take an optional parameter for the number of eleme...

lib/redis/connection/memory.rb

@@ -473,10 +473,8 @@ def sdiffstore(destination, key1, *keys)
         data[destination] = ::Set.new(result)
       end
 
-      def srandmember(key)
-        data_type_check(key, ::Set)
-        return nil unless data[key]
-        data[key].to_a[rand(data[key].size)]
+      def srandmember(key, number=nil)
+        number.nil? ? srandmember_single(key) : srandmember_multiple(key, number)
       end
 
       def del(*keys)
@@ -975,6 +973,25 @@ def get_limit(opts, vals)
         def mapped_param? param
           param.size == 1 && param[0].is_a?(Array)
         end
+
+        def srandmember_single(key)
+          data_type_check(key, ::Set)
+          return nil unless data[key]
+          data[key].to_a[rand(data[key].size)]
+        end
+
+        def srandmember_multiple(key, number)
+          return [] unless data[key]
+          if number >= 0
+            # replace with `data[key].to_a.sample(number)` when 1.8.7 is deprecated
+            (1..number).inject([]) do |selected, _|
+              available_elements = data[key].to_a - selected
+              selected << available_elements[rand(available_elements.size)]
+            end.compact
+          else
+            (1..-number).map { data[key].to_a[rand(data[key].size)] }.flatten
+          end
+        end
     end
   end
 end

spec/sets_spec.rb

@@ -182,4 +182,80 @@ module FakeRedis
       @client.smembers("key").should =~ ["a", "b", "c", "d", "e"]
     end
   end
+
+  describe 'srandmember' do
+    before(:each) do
+      @client = Redis.new
+    end
+
+    context 'with a set that has three elements' do
+      before do
+        @client.sadd("key1", "a")
+        @client.sadd("key1", "b")
+        @client.sadd("key1", "c")
+      end
+
+      context 'when called without the optional number parameter' do
+        it 'is a random element from the set' do
+          random_element = @client.srandmember("key1")
+
+          ['a', 'b', 'c'].include?(random_element).should be_true
+        end
+      end
+
+      context 'when called with the optional number parameter of 1' do
+        it 'is an array of one random element from the set' do
+          random_elements = @client.srandmember("key1", 1)
+
+          [['a'], ['b'], ['c']].include?(@client.srandmember("key1", 1)).should be_true
+        end
+      end
+
+      context 'when called with the optional number parameter of 2' do
+        it 'is an array of two unique, random elements from the set' do
+          random_elements = @client.srandmember("key1", 2)
+
+          random_elements.count.should == 2
+          random_elements.uniq.count.should == 2
+          random_elements.all? do |element|
+            ['a', 'b', 'c'].include?(element).should be_true
+          end
+        end
+      end
+
+      context 'when called with an optional parameter of -100' do
+        it 'is an array of 100 random elements from the set, some of which are repeated' do
+          random_elements = @client.srandmember("key1", -100)
+
+          random_elements.count.should == 100
+          random_elements.uniq.count.should <= 3
+          random_elements.all? do |element|
+            ['a', 'b', 'c'].include?(element).should be_true
+          end
+        end
+      end
+
+      context 'when called with an optional parameter of 100' do
+        it 'is an array of all of the elements from the set, none of which are repeated' do
+          random_elements = @client.srandmember("key1", 100)
+
+          random_elements.count.should == 3
+          random_elements.uniq.count.should == 3
+          random_elements.should =~ ['a', 'b', 'c']
+        end
+      end
+    end
+
+    context 'with an empty set' do
+      before { @client.del("key1") }
+
+      it 'is nil without the extra parameter' do
+        @client.srandmember("key1").should be_nil
+      end
+
+      it 'is an empty array with an extra parameter' do
+        @client.srandmember("key1", 1).should == []
+      end
+    end
+  end
 end