From 24f7ed04996f4ed51f43dfbebe41b5c88ac7b6fa Mon Sep 17 00:00:00 2001
From: diamondburned <datutbrus@gmail.com>
Date: Sat, 18 Jul 2020 18:33:07 -0700
Subject: [PATCH] Gateway: ReconnectCtx now returns error; fixed test

---
 gateway/gateway.go          | 14 ++++++++++----
 gateway/integration_test.go |  8 ++++++--
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/gateway/gateway.go b/gateway/gateway.go
index a2b1dd4..7700708 100644
--- a/gateway/gateway.go
+++ b/gateway/gateway.go
@@ -203,12 +203,18 @@ func (g *Gateway) Close() error {
 // Reconnect tries to reconnect forever. It will resume the connection if
 // possible. If an Invalid Session is received, it will start a fresh one.
 func (g *Gateway) Reconnect() {
-	g.ReconnectCtx(context.Background())
+	for {
+		if err := g.ReconnectCtx(context.Background()); err != nil {
+			g.ErrorLog(err)
+		} else {
+			return
+		}
+	}
 }
 
 // ReconnectCtx attempts to reconnect until context expires. If context cannot
 // expire, then the gateway will try to reconnect forever.
-func (g *Gateway) ReconnectCtx(ctx context.Context) {
+func (g *Gateway) ReconnectCtx(ctx context.Context) error {
 	wsutil.WSDebug("Reconnecting...")
 
 	// Guarantee the gateway is already closed. Ignore its error, as we're
@@ -223,11 +229,11 @@ func (g *Gateway) ReconnectCtx(ctx context.Context) {
 		// https://discordapp.com/developers/docs/topics/gateway#rate-limiting
 
 		if err := g.OpenContext(ctx); err != nil {
-			g.ErrorLog(errors.Wrap(err, "failed to open gateway"))
+			return errors.Wrap(err, "failed to open gateway")
 		}
 
 		wsutil.WSDebug("Started after attempt:", i)
-		return
+		return nil
 	}
 }
 
diff --git a/gateway/integration_test.go b/gateway/integration_test.go
index 5e05aee..04645fe 100644
--- a/gateway/integration_test.go
+++ b/gateway/integration_test.go
@@ -3,6 +3,7 @@
 package gateway
 
 import (
+	"context"
 	"log"
 	"os"
 	"strings"
@@ -78,9 +79,12 @@ func TestIntegration(t *testing.T) {
 	// Sleep past the rate limiter before reconnecting:
 	time.Sleep(5 * time.Second)
 
-	// Try and reconnect forever:
 	gotimeout(t, func() {
-		if err := gateway.Reconnect(); err != nil {
+		// Try and reconnect for 20 seconds maximum.
+		ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+		defer cancel()
+
+		if err := gateway.ReconnectCtx(ctx); err != nil {
 			t.Fatal("Unexpected error while reconnecting:", err)
 		}
 	})