Add support for optional count arg in rpop

lib/mock_redis/list_methods.rb

@@ -186,8 +186,13 @@ def ltrim(key, start, stop)
       end
     end
 
-    def rpop(key)
-      with_list_at(key) { |list| list&.pop }
+    def rpop(key, count = nil)
+      return with_list_at(key) { |list| list&.pop } if count.nil?
+
+      record_count = llen(key)
+      return nil if record_count.zero?
+
+      [record_count, count].min.times.map { with_list_at(key, &:pop) }
     end
 
     def rpoplpush(source, destination)

spec/commands/rpop_spec.rb

@@ -30,6 +30,37 @@
     expect(@redises.get(@key)).to be_nil
   end
 
-  let(:default_error) { RedisMultiplexer::MismatchedResponse }
-  it_should_behave_like 'a list-only command'
+  context 'when count != nil' do
+    it 'returns array with one element if count == 1' do
+      @redises.rpush(@key, %w[one two three four five])
+
+      expect(@redises.rpop(@key, 1)).to eq(%w[five])
+      expect(@redises.lrange(@key, 0, -1)).to eq(%w[one two three four])
+    end
+
+    it 'returns the number of records requested' do
+      @redises.rpush(@key, %w[one two three four five])
+
+      expect(@redises.rpop(@key, 2)).to eq(%w[five four])
+      expect(@redises.lrange(@key, 0, -1)).to eq(%w[one two three])
+    end
+
+    it 'returns nil for nonexistent key' do
+      expect(@redises.rpop(@key, 2)).to be_nil
+    end
+
+    it 'returns all records when requesting more than list length' do
+      @redises.rpush(@key, %w[one two three])
+
+      expect(@redises.rpop(@key, 10)).to eq(%w[three two one])
+      expect(@redises.lrange(@key, 0, -1)).to eq([])
+    end
+  end
+
+  it_should_behave_like 'a list-only command' do
+    let(:args) { [1] }
+    let(:error) do
+      [Redis::CommandError, 'WRONGTYPE Operation against a key holding the wrong kind of value']
+    end
+  end
 end

spec/support/shared_examples/does_not_cleanup_empty_strings.rb

@@ -1,14 +1,16 @@
 RSpec.shared_examples_for 'does not remove empty strings on error' do
-  it 'does not remove empty strings on error' do |example|
-    key = 'mock-redis-test:not-a-string'
+  let(:method) { |example| method_from_description(example) }
+  let(:args) { args_for_method(method) }
+  let(:error) { defined?(default_error) ? default_error : RuntimeError }
 
-    method = method_from_description(example)
-    args = args_for_method(method).unshift(key)
+  it 'does not remove empty strings on error' do
+    key = 'mock-redis-test:not-a-string'
+    key_and_args = args.unshift(key)
 
     @redises.set(key, '')
     expect do
-      @redises.send(method, *args)
-    end.to raise_error(defined?(default_error) ? default_error : RuntimeError)
+      @redises.send(method, *key_and_args)
+    end.to raise_error(*error)
     expect(@redises.get(key)).to eq('')
   end
 end

spec/support/shared_examples/only_operates_on_lists.rb

@@ -1,15 +1,18 @@
 RSpec.shared_examples_for 'a list-only command' do
-  it 'raises an error for non-list values' do |example|
-    key = 'mock-redis-test:list-only'
+  let(:method) { |example| method_from_description(example) }
+  let(:args) { args_for_method(method) }
+  let(:error) { defined?(default_error) ? default_error : RuntimeError }
 
-    method = method_from_description(example)
-    args = args_for_method(method).unshift(key)
+  it 'raises an error for non-list values' do
+    key = 'mock-redis-test:list-only'
+    key_and_args = args.unshift(key)
 
     @redises.set(key, 1)
+
     expect do
-      @redises.send(method, *args)
-    end.to raise_error(defined?(default_error) ? default_error : RuntimeError)
+      @redises.send(method, *key_and_args)
+    end.to raise_error(*error)
   end
 
-  it_should_behave_like 'does not remove empty strings on error'
+  include_examples 'does not remove empty strings on error'
 end