From bd8aa68cf13e2866166baa1359a7d809d233bd13 Mon Sep 17 00:00:00 2001
From: "STeve (Xin) Huang" <xin.huang@goteleport.com>
Date: Mon, 27 Nov 2023 16:11:28 -0500
Subject: [PATCH] Enable StatusString for all readers (#50)

---
 internal/proto/reader.go      | 26 ++++----------------------
 internal/proto/writer_test.go |  2 +-
 proto.go                      |  2 +-
 pubsub.go                     |  4 ++++
 redis_test.go                 |  3 ++-
 5 files changed, 12 insertions(+), 25 deletions(-)

diff --git a/internal/proto/reader.go b/internal/proto/reader.go
index d1a833ffa..fe8ec7415 100644
--- a/internal/proto/reader.go
+++ b/internal/proto/reader.go
@@ -55,28 +55,13 @@ func ParseErrorReply(line []byte) error {
 
 //------------------------------------------------------------------------------
 
-type readerOpt struct {
-	// useStatusStringType uses `StatusString` type instead of string for RespStatus.
-	useStatusStringType bool
-}
-
-func ReaderOptUseStatusStringType(o *readerOpt) {
-	o.useStatusStringType = true
-}
-
 type Reader struct {
-	rd  *bufio.Reader
-	opt readerOpt
+	rd *bufio.Reader
 }
 
-func NewReader(rd io.Reader, opts ...func(o *readerOpt)) *Reader {
-	opt := readerOpt{}
-	for _, applyOpt := range opts {
-		applyOpt(&opt)
-	}
+func NewReader(rd io.Reader) *Reader {
 	return &Reader{
-		rd:  bufio.NewReader(rd),
-		opt: opt,
+		rd: bufio.NewReader(rd),
 	}
 }
 
@@ -177,10 +162,7 @@ func (r *Reader) ReadReply() (interface{}, error) {
 
 	switch line[0] {
 	case RespStatus:
-		if r.opt.useStatusStringType {
-			return StatusString(line[1:]), nil
-		}
-		return string(line[1:]), nil
+		return StatusString(line[1:]), nil
 	case RespInt:
 		return util.ParseInt(line[1:], 10, 64)
 	case RespFloat:
diff --git a/internal/proto/writer_test.go b/internal/proto/writer_test.go
index 865f89eb1..ff7578ca5 100644
--- a/internal/proto/writer_test.go
+++ b/internal/proto/writer_test.go
@@ -105,7 +105,7 @@ func TestWriteStatus(t *testing.T) {
 	inputStatusBytes := []byte("+status\r\n")
 
 	// Read it.
-	reader := proto.NewReader(bytes.NewReader(inputStatusBytes), proto.ReaderOptUseStatusStringType)
+	reader := proto.NewReader(bytes.NewReader(inputStatusBytes))
 	readStatus, err := reader.ReadReply()
 	if err != nil {
 		t.Errorf("Failed to ReadReply: %v", err)
diff --git a/proto.go b/proto.go
index 2b0772297..d942290c7 100644
--- a/proto.go
+++ b/proto.go
@@ -15,7 +15,7 @@ const RespArray = proto.RespArray
 const RespInt = proto.RespInt
 
 func NewReader(rd io.Reader) *Reader {
-	return proto.NewReader(rd, proto.ReaderOptUseStatusStringType)
+	return proto.NewReader(rd)
 }
 
 func NewWriter(wr *bytes.Buffer) *Writer {
diff --git a/pubsub.go b/pubsub.go
index 16c0f5672..596b6878e 100644
--- a/pubsub.go
+++ b/pubsub.go
@@ -366,6 +366,10 @@ func (p *Pong) String() string {
 
 func (c *PubSub) newMessage(reply interface{}) (interface{}, error) {
 	switch reply := reply.(type) {
+	case proto.StatusString:
+		return &Pong{
+			Payload: string(reply),
+		}, nil
 	case string:
 		return &Pong{
 			Payload: reply,
diff --git a/redis_test.go b/redis_test.go
index 6d3842070..39b6affa5 100644
--- a/redis_test.go
+++ b/redis_test.go
@@ -12,6 +12,7 @@ import (
 	. "github.com/bsm/gomega"
 
 	"github.com/redis/go-redis/v9"
+	"github.com/redis/go-redis/v9/internal/proto"
 )
 
 type redisHookError struct{}
@@ -192,7 +193,7 @@ var _ = Describe("Client", func() {
 		Expect(client.Echo(ctx, "hello").Err()).NotTo(HaveOccurred())
 
 		Expect(cmd.Err()).NotTo(HaveOccurred())
-		Expect(cmd.Val()).To(Equal("PONG"))
+		Expect(cmd.Val()).To(Equal(proto.StatusString("PONG")))
 	})
 
 	It("should retry command on network error", func() {