From a54ed5afb23fa9f12b312fa9624768487d97b6dc Mon Sep 17 00:00:00 2001
From: Dom Parfitt <mail@domparfitt.com>
Date: Thu, 3 Nov 2022 09:22:54 +0000
Subject: [PATCH] test: add withConn tests

---
 with_conn_test.go | 81 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 with_conn_test.go

diff --git a/with_conn_test.go b/with_conn_test.go
new file mode 100644
index 00000000..4c90d650
--- /dev/null
+++ b/with_conn_test.go
@@ -0,0 +1,81 @@
+package redis
+
+import (
+	"context"
+
+	"github.com/go-redis/redis/v9/internal/pool"
+	"github.com/go-redis/redis/v9/internal/proto"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+type timeoutErr struct {
+	error
+}
+
+func (e timeoutErr) Timeout() bool {
+	return true
+}
+
+func (e timeoutErr) Temporary() bool {
+	return true
+}
+
+func (e timeoutErr) Error() string {
+	return "i/o timeout"
+}
+
+var _ = Describe("withConn", func() {
+	var client *Client
+
+	BeforeEach(func() {
+		client = NewClient(&Options{
+			PoolSize: 1,
+		})
+	})
+
+	AfterEach(func() {
+		client.Close()
+	})
+
+	It("should replace the connection in the pool when there is no error", func() {
+		var conn *pool.Conn
+
+		client.withConn(ctx, func(ctx context.Context, c *pool.Conn) error {
+			conn = c
+			return nil
+		})
+
+		newConn, err := client.connPool.Get(ctx)
+		Expect(err).To(BeNil())
+		Expect(newConn).To(Equal(conn))
+	})
+
+	It("should replace the connection in the pool when there is an error not related to a bad connection", func() {
+		var conn *pool.Conn
+
+		client.withConn(ctx, func(ctx context.Context, c *pool.Conn) error {
+			conn = c
+			return proto.RedisError("LOADING")
+		})
+
+		newConn, err := client.connPool.Get(ctx)
+		Expect(err).To(BeNil())
+		Expect(newConn).To(Equal(conn))
+	})
+
+	It("should remove the connection from the pool when it times out", func() {
+		var conn *pool.Conn
+
+		client.withConn(ctx, func(ctx context.Context, c *pool.Conn) error {
+			conn = c
+			return timeoutErr{}
+		})
+
+		newConn, err := client.connPool.Get(ctx)
+		Expect(err).To(BeNil())
+		Expect(newConn).NotTo(Equal(conn))
+		Expect(client.connPool.Len()).To(Equal(1))
+	})
+})