Fix WithContext followed by WrapProcess

redis_test.go

@@ -228,18 +228,40 @@ var _ = Describe("Client", func() {
 	})
 
 	It("should call WrapProcess", func() {
-		var wrapperFnCalled bool
+		var fnCalled bool
 
-		client.WrapProcess(func(oldProcess func(redis.Cmder) error) func(redis.Cmder) error {
+		client.WrapProcess(func(old func(redis.Cmder) error) func(redis.Cmder) error {
 			return func(cmd redis.Cmder) error {
-				wrapperFnCalled = true
-				return oldProcess(cmd)
+				fnCalled = true
+				return old(cmd)
 			}
 		})
 
 		Expect(client.Ping().Err()).NotTo(HaveOccurred())
+		Expect(fnCalled).To(BeTrue())
+	})
 
-		Expect(wrapperFnCalled).To(BeTrue())
+	It("should call WrapProcess after WithContext", func() {
+		var fn1Called, fn2Called bool
+
+		client.WrapProcess(func(old func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+			return func(cmd redis.Cmder) error {
+				fn1Called = true
+				return old(cmd)
+			}
+		})
+
+		client2 := client.WithContext(client.Context())
+		client2.WrapProcess(func(old func(cmd redis.Cmder) error) func(cmd redis.Cmder) error {
+			return func(cmd redis.Cmder) error {
+				fn2Called = true
+				return old(cmd)
+			}
+		})
+
+		Expect(client2.Ping().Err()).NotTo(HaveOccurred())
+		Expect(fn2Called).To(BeTrue())
+		Expect(fn1Called).To(BeTrue())
 	})
 })