From ab5331e8f02863c1ff4c02612073a90b0360632a Mon Sep 17 00:00:00 2001
From: roertbb <robert.banaszak1997@gmail.com>
Date: Fri, 20 Oct 2023 15:37:38 +0200
Subject: [PATCH] fix(commands): add count argument to lpop and remove key when
 list is empty

---
 src/commands/lpop.js              | 31 +++++++++++++++----
 test/integration/commands/lpop.js | 50 +++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/src/commands/lpop.js b/src/commands/lpop.js
index d6a6c43e2..267d2d3e3 100644
--- a/src/commands/lpop.js
+++ b/src/commands/lpop.js
@@ -1,19 +1,40 @@
 import { convertStringToBuffer } from '../commands-utils/convertStringToBuffer'
 
-export function lpop(key) {
+export function lpop(key, count) {
   if (this.data.has(key) && !(this.data.get(key) instanceof Array)) {
     throw new Error(`Key ${key} does not contain a list`)
   }
   const list = this.data.get(key) || []
 
-  const item = list.length > 0 ? list.shift() : null
+  if (list.length === 0) {
+    return null
+  }
+
+  if (count) {
+    const items = list.slice(0, count)
+    const remainingItems = list.slice(count)
 
-  this.data.set(key, list)
+    if (remainingItems.length > 0) {
+      this.data.set(key, remainingItems)
+    } else {
+      this.data.delete(key)
+    }
+
+    return items
+  }
+
+  const item = list.shift()
+
+  if (list.length > 0) {
+    this.data.set(key, list)
+  } else {
+    this.data.delete(key)
+  }
 
   return item
 }
 
-export function lpopBuffer(key) {
-  const val = lpop.apply(this, [key])
+export function lpopBuffer(key, count) {
+  const val = lpop.apply(this, [key, count])
   return convertStringToBuffer(val)
 }
diff --git a/test/integration/commands/lpop.js b/test/integration/commands/lpop.js
index 087876a90..a49adbcbc 100644
--- a/test/integration/commands/lpop.js
+++ b/test/integration/commands/lpop.js
@@ -21,6 +21,56 @@ runTwinSuite('lpop', (command, equals) => {
       }
     )
 
+    // @TODO Rewrite test so it runs on a real Redis instance
+    ;(process.env.IS_E2E ? it.skip : it)(
+      'should remove key and return first elements when it was last on the list',
+      () => {
+        const redis = new Redis({
+          data: {
+            foo: ['1'],
+          },
+        })
+
+        return redis[command]('foo')
+          .then(result => expect(equals(result, '1')).toBe(true))
+          .then(() => redis.exists('foo'))
+          .then(status => expect(status).toBe(0))
+      }
+    )
+
+    // @TODO Rewrite test so it runs on a real Redis instance
+    ;(process.env.IS_E2E ? it.skip : it)(
+      'should remove and return "count" elements if second argument is provided',
+      () => {
+        const redis = new Redis({
+          data: {
+            foo: ['5', '4', '3', '2', '1'],
+          },
+        })
+
+        return redis[command]('foo', 2)
+          .then(result => expect(result.map(v => Buffer.isBuffer(v) ? v.toString() : v)).toEqual(['5', '4']))
+          .then(() => expect(redis.data.get('foo')).toEqual(['3', '2', '1']))
+      }
+    )
+
+    // @TODO Rewrite test so it runs on a real Redis instance
+    ;(process.env.IS_E2E ? it.skip : it)(
+      'should remove key and return all elements on larger number if second argument is provided',
+      () => {
+        const redis = new Redis({
+          data: {
+            foo: ['5', '4', '3', '2', '1'],
+          },
+        })
+
+        return redis[command]('foo', 7)
+          .then(result => expect(result.map(v => Buffer.isBuffer(v) ? v.toString() : v)).toEqual(['5', '4', '3', '2', '1']))
+          .then(() => redis.exists('foo'))
+          .then(status => expect(status).toBe(0))
+      }
+    )
+
     // @TODO Rewrite test so it runs on a real Redis instance
     ;(process.env.IS_E2E ? it.skip : it)(
       'should return buffer values correctly as buffer',