From 00de46e290cf0a47bb78a9e8aa264f4b6ce941cd Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Fri, 27 Sep 2019 16:21:44 -0700
Subject: [PATCH] fix(test): fix a flaky pubsub test

---
 tests/pubsub.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/tests/pubsub.go b/tests/pubsub.go
index e662915..36353f8 100644
--- a/tests/pubsub.go
+++ b/tests/pubsub.go
@@ -34,13 +34,20 @@ func (tp *TestSuite) TestBasicPubSub(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	done := make(chan struct{})
 	go func() {
+		defer close(done)
+
 		ticker := time.NewTicker(100 * time.Millisecond)
 		defer ticker.Stop()
 
 		for {
 			err := apis[1].PubSub().Publish(ctx, "testch", []byte("hello world"))
-			if err != nil {
+			switch err {
+			case nil:
+			case context.Canceled:
+				return
+			default:
 				t.Error(err)
 				cancel()
 				return
@@ -53,6 +60,13 @@ func (tp *TestSuite) TestBasicPubSub(t *testing.T) {
 		}
 	}()
 
+	// Wait for the sender to finish before we return.
+	// Otherwise, we can get random errors as publish fails.
+	defer func() {
+		cancel()
+		<-done
+	}()
+
 	m, err := sub.Next(ctx)
 	if err != nil {
 		t.Fatal(err)