fix(test): fix a flaky pubsub test

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)