From cb7f3e8a17c6ec6f7f5427ed05d3cca631ee1844 Mon Sep 17 00:00:00 2001
From: fabio <zzzfwww@163.com>
Date: Thu, 9 Mar 2023 13:45:45 +0800
Subject: [PATCH] feat(redis):add LpopCount,RpopCount (#2990)

---
 core/stores/redis/redis.go      | 40 +++++++++++++++++++++++++++++++++
 core/stores/redis/redis_test.go | 36 +++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/core/stores/redis/redis.go b/core/stores/redis/redis.go
index 8004986b..31c7e514 100644
--- a/core/stores/redis/redis.go
+++ b/core/stores/redis/redis.go
@@ -1170,6 +1170,26 @@ func (s *Redis) LpopCtx(ctx context.Context, key string) (val string, err error)
 	return
 }
 
+// LpopCount is the implementation of redis lpopCount command.
+func (s *Redis) LpopCount(key string, count int) ([]string, error) {
+	return s.LpopCountCtx(context.Background(), key, count)
+}
+
+// LpopCountCtx is the implementation of redis lpopCount command.
+func (s *Redis) LpopCountCtx(ctx context.Context, key string, count int) (val []string, err error) {
+	err = s.brk.DoWithAcceptable(func() error {
+		conn, err := getRedis(s)
+		if err != nil {
+			return err
+		}
+
+		val, err = conn.LPopCount(ctx, key, count).Result()
+		return err
+	}, acceptable)
+
+	return
+}
+
 // Lpush is the implementation of redis lpush command.
 func (s *Redis) Lpush(key string, values ...any) (int, error) {
 	return s.LpushCtx(context.Background(), key, values...)
@@ -1432,6 +1452,26 @@ func (s *Redis) RpopCtx(ctx context.Context, key string) (val string, err error)
 	return
 }
 
+// RpopCount is the implementation of redis rpopCount command.
+func (s *Redis) RpopCount(key string, count int) ([]string, error) {
+	return s.RpopCountCtx(context.Background(), key, count)
+}
+
+// RpopCountCtx is the implementation of redis rpopCount command.
+func (s *Redis) RpopCountCtx(ctx context.Context, key string, count int) (val []string, err error) {
+	err = s.brk.DoWithAcceptable(func() error {
+		conn, err := getRedis(s)
+		if err != nil {
+			return err
+		}
+
+		val, err = conn.RPopCount(ctx, key, count).Result()
+		return err
+	}, acceptable)
+
+	return
+}
+
 // Rpush is the implementation of redis rpush command.
 func (s *Redis) Rpush(key string, values ...any) (int, error) {
 	return s.RpushCtx(context.Background(), key, values...)
diff --git a/core/stores/redis/redis_test.go b/core/stores/redis/redis_test.go
index 88ddd48b..ffd4d00d 100644
--- a/core/stores/redis/redis_test.go
+++ b/core/stores/redis/redis_test.go
@@ -507,6 +507,14 @@ func TestRedis_List(t *testing.T) {
 			vals, err = client.Lrange("key", 0, 10)
 			assert.Nil(t, err)
 			assert.EqualValues(t, []string{"value2", "value3"}, vals)
+			vals, err = client.LpopCount("key", 2)
+			assert.Nil(t, err)
+			assert.EqualValues(t, []string{"value2", "value3"}, vals)
+			_, err = client.Lpush("key", "value1", "value2")
+			assert.Nil(t, err)
+			vals, err = client.RpopCount("key", 4)
+			assert.Nil(t, err)
+			assert.EqualValues(t, []string{"value1", "value2"}, vals)
 		})
 	})
 
@@ -523,6 +531,34 @@ func TestRedis_List(t *testing.T) {
 
 			_, err = client.Rpush("key", "value3", "value4")
 			assert.Error(t, err)
+
+			_, err = client.LpopCount("key", 2)
+			assert.Error(t, err)
+
+			_, err = client.RpopCount("key", 2)
+			assert.Error(t, err)
+		})
+	})
+	t.Run("list redis type error", func(t *testing.T) {
+		runOnRedisWithError(t, func(client *Redis) {
+			client.Type = "nil"
+			_, err := client.Llen("key")
+			assert.Error(t, err)
+
+			_, err = client.Lpush("key", "value1", "value2")
+			assert.Error(t, err)
+
+			_, err = client.Lrem("key", 2, "value1")
+			assert.Error(t, err)
+
+			_, err = client.Rpush("key", "value3", "value4")
+			assert.Error(t, err)
+
+			_, err = client.LpopCount("key", 2)
+			assert.Error(t, err)
+
+			_, err = client.RpopCount("key", 2)
+			assert.Error(t, err)
 		})
 	})
 }