From 787609d1e6860f0a84bf274ad1bf5bab2b8fd856 Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Sat, 9 Apr 2016 11:34:40 +0300
Subject: [PATCH] Fix race in tests.

---
 export_test.go | 4 ++--
 pubsub.go      | 8 +++++---
 pubsub_test.go | 5 ++---
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/export_test.go b/export_test.go
index f071b54..587f2e5 100644
--- a/export_test.go
+++ b/export_test.go
@@ -14,6 +14,6 @@ func (c *PubSub) Pool() pool.Pooler {
 	return c.base.connPool
 }
 
-func SetReceiveMessageTimeout(d time.Duration) {
-	receiveMessageTimeout = d
+func (c *PubSub) ReceiveMessageTimeout(timeout time.Duration) (*Message, error) {
+	return c.receiveMessage(timeout)
 }
diff --git a/pubsub.go b/pubsub.go
index 1bf2f93..20cf1fb 100644
--- a/pubsub.go
+++ b/pubsub.go
@@ -8,8 +8,6 @@ import (
 	"gopkg.in/redis.v3/internal/pool"
 )
 
-var receiveMessageTimeout = 5 * time.Second
-
 // Posts a message to the given channel.
 func (c *Client) Publish(channel, message string) *IntCmd {
 	req := NewIntCmd("PUBLISH", channel, message)
@@ -255,9 +253,13 @@ func (c *PubSub) Receive() (interface{}, error) {
 // messages. It automatically reconnects to Redis Server and resubscribes
 // to channels in case of network errors.
 func (c *PubSub) ReceiveMessage() (*Message, error) {
+	return c.receiveMessage(5 * time.Second)
+}
+
+func (c *PubSub) receiveMessage(timeout time.Duration) (*Message, error) {
 	var errNum uint
 	for {
-		msgi, err := c.ReceiveTimeout(receiveMessageTimeout)
+		msgi, err := c.ReceiveTimeout(timeout)
 		if err != nil {
 			if !isNetworkError(err) {
 				return nil, err
diff --git a/pubsub_test.go b/pubsub_test.go
index df0dd94..aa0e465 100644
--- a/pubsub_test.go
+++ b/pubsub_test.go
@@ -256,8 +256,7 @@ var _ = Describe("PubSub", func() {
 	})
 
 	It("should ReceiveMessage after timeout", func() {
-		timeout := time.Second
-		redis.SetReceiveMessageTimeout(timeout)
+		timeout := 100 * time.Millisecond
 
 		pubsub, err := client.Subscribe("mychannel")
 		Expect(err).NotTo(HaveOccurred())
@@ -276,7 +275,7 @@ var _ = Describe("PubSub", func() {
 			Expect(n).To(Equal(int64(1)))
 		}()
 
-		msg, err := pubsub.ReceiveMessage()
+		msg, err := pubsub.ReceiveMessageTimeout(timeout)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(msg.Channel).To(Equal("mychannel"))
 		Expect(msg.Payload).To(Equal("hello"))