From 5eee83fd726bae67e411a6290a1fd7a3edabfffe Mon Sep 17 00:00:00 2001
From: rhymes <rhymes@hey.com>
Date: Sun, 21 Mar 2021 11:12:22 +0100
Subject: [PATCH] [Redis 6.2] Add count argument to lpop and rpop

---
 lib/redis.rb             | 22 ++++++++++++++--------
 lib/redis/distributed.rb | 12 ++++++------
 test/lint/lists.rb       | 22 ++++++++++++++++++++++
 3 files changed, 42 insertions(+), 14 deletions(-)

diff --git a/lib/redis.rb b/lib/redis.rb
index fc3a47cea..19aef5374 100644
--- a/lib/redis.rb
+++ b/lib/redis.rb
@@ -1172,23 +1172,29 @@ def rpushx(key, value)
     end
   end
 
-  # Remove and get the first element in a list.
+  # Remove and get the first elements in a list.
   #
   # @param [String] key
-  # @return [String]
-  def lpop(key)
+  # @param [Integer] count number of elements to remove
+  # @return [String, Array<String>] the values of the first elements
+  def lpop(key, count = nil)
     synchronize do |client|
-      client.call([:lpop, key])
+      command = [:lpop, key]
+      command << count if count
+      client.call(command)
     end
   end
 
-  # Remove and get the last element in a list.
+  # Remove and get the last elements in a list.
   #
   # @param [String] key
-  # @return [String]
-  def rpop(key)
+  # @param [Integer] count number of elements to remove
+  # @return [String, Array<String>] the values of the last elements
+  def rpop(key, count = nil)
     synchronize do |client|
-      client.call([:rpop, key])
+      command = [:rpop, key]
+      command << count if count
+      client.call(command)
     end
   end
 
diff --git a/lib/redis/distributed.rb b/lib/redis/distributed.rb
index 0ffaa2bd2..6d1566550 100644
--- a/lib/redis/distributed.rb
+++ b/lib/redis/distributed.rb
@@ -413,14 +413,14 @@ def rpushx(key, value)
       node_for(key).rpushx(key, value)
     end
 
-    # Remove and get the first element in a list.
-    def lpop(key)
-      node_for(key).lpop(key)
+    # Remove and get the first elements in a list.
+    def lpop(key, count = nil)
+      node_for(key).lpop(key, count)
     end
 
-    # Remove and get the last element in a list.
-    def rpop(key)
-      node_for(key).rpop(key)
+    # Remove and get the last elements in a list.
+    def rpop(key, count = nil)
+      node_for(key).rpop(key, count)
     end
 
     # Remove the last element in a list, append it to another list and return
diff --git a/test/lint/lists.rb b/test/lint/lists.rb
index b0c5a2570..e8efea07c 100644
--- a/test/lint/lists.rb
+++ b/test/lint/lists.rb
@@ -119,6 +119,17 @@ def test_lpop
       assert_equal 1, r.llen("foo")
     end
 
+    def test_lpop_count
+      target_version("6.2") do
+        r.rpush "foo", "s1"
+        r.rpush "foo", "s2"
+
+        assert_equal 2, r.llen("foo")
+        assert_equal ["s1", "s2"], r.lpop("foo", 2)
+        assert_equal 0, r.llen("foo")
+      end
+    end
+
     def test_rpop
       r.rpush "foo", "s1"
       r.rpush "foo", "s2"
@@ -128,6 +139,17 @@ def test_rpop
       assert_equal 1, r.llen("foo")
     end
 
+    def test_rpop_count
+      target_version("6.2") do
+        r.rpush "foo", "s1"
+        r.rpush "foo", "s2"
+
+        assert_equal 2, r.llen("foo")
+        assert_equal ["s2", "s1"], r.rpop("foo", 2)
+        assert_equal 0, r.llen("foo")
+      end
+    end
+
     def test_linsert
       r.rpush "foo", "s1"
       r.rpush "foo", "s3"