From 341966d63a69649a3fabe288b633c63139a17dd2 Mon Sep 17 00:00:00 2001
From: Eric Chlebek <eric@sensu.io>
Date: Wed, 14 Feb 2024 11:15:48 -0800
Subject: [PATCH] Improve coverage, revert zero duration change

---
 client/wsclient.go      |  6 +-----
 client/wsclient_test.go | 28 +++++++++++++++++++++++-----
 internal/retryafter.go  |  6 ------
 3 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/client/wsclient.go b/client/wsclient.go
index c2e9237c..8c725dc5 100644
--- a/client/wsclient.go
+++ b/client/wsclient.go
@@ -3,7 +3,6 @@ package client
 import (
 	"context"
 	"errors"
-	"fmt"
 	"net/http"
 	"net/url"
 	"sync"
@@ -137,7 +136,7 @@ func (c *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sh
 				// very liberal handling of 3xx that largely ignores HTTP semantics
 				redirect, err := resp.Location()
 				if err != nil {
-					fmt.Println("error in location", err)
+					c.common.Callbacks.OnConnectFailed(ctx, err)
 					c.common.Logger.Errorf(ctx, "%d redirect, but no valid location: %s", resp.StatusCode, err)
 					return err, duration
 				}
@@ -151,9 +150,6 @@ func (c *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sh
 				// Set the URL to the redirect, so that it connects to it on the
 				// next cycle.
 				c.url = redirect
-
-				// don't delay in following the redirect
-				duration = sharedinternal.ZeroDuration
 			} else {
 				c.common.Logger.Errorf(ctx, "Server responded with status=%v", resp.Status)
 			}
diff --git a/client/wsclient_test.go b/client/wsclient_test.go
index 4e87a3b0..f828e059 100644
--- a/client/wsclient_test.go
+++ b/client/wsclient_test.go
@@ -181,25 +181,37 @@ func TestVerifyWSCompress(t *testing.T) {
 	}
 }
 
-func redirectServer(to string) *httptest.Server {
+func redirectServer(to string, status int) *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		http.Redirect(w, req, to, http.StatusSeeOther)
 	}))
 }
 
+func errServer() *httptest.Server {
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		w.WriteHeader(302)
+	}))
+}
+
 func TestRedirectWS(t *testing.T) {
 	redirectee := internal.StartMockServer(t)
 	tests := []struct {
 		Name       string
 		Redirector *httptest.Server
+		ExpError   bool
 	}{
 		{
 			Name:       "redirect ws scheme",
-			Redirector: redirectServer("ws://" + redirectee.Endpoint),
+			Redirector: redirectServer("ws://"+redirectee.Endpoint, 302),
 		},
 		{
 			Name:       "redirect http scheme",
-			Redirector: redirectServer("http://" + redirectee.Endpoint),
+			Redirector: redirectServer("http://"+redirectee.Endpoint, 302),
+		},
+		{
+			Name:       "missing location header",
+			Redirector: errServer(),
+			ExpError:   true,
 		},
 	}
 
@@ -233,8 +245,14 @@ func TestRedirectWS(t *testing.T) {
 			startClient(t, settings, client)
 
 			// Wait for connection to be established.
-			eventually(t, func() bool { return conn.Load() != nil })
-			assert.True(t, connectErr.Load() == nil)
+			eventually(t, func() bool { return conn.Load() != nil || connectErr.Load() != nil })
+			if test.ExpError {
+				if connectErr.Load() == nil {
+					t.Error("expected non-nil error")
+				}
+			} else {
+				assert.True(t, connectErr.Load() == nil)
+			}
 
 			// Stop the client.
 			err = client.Stop(context.Background())
diff --git a/internal/retryafter.go b/internal/retryafter.go
index 25f168c7..784473bd 100644
--- a/internal/retryafter.go
+++ b/internal/retryafter.go
@@ -17,12 +17,6 @@ type OptionalDuration struct {
 	Defined bool
 }
 
-// ZeroDuration represents a zero length duration that is defined.
-var ZeroDuration = OptionalDuration{
-	Duration: 0,
-	Defined:  true,
-}
-
 func parseDelaySeconds(s string) (time.Duration, error) {
 	n, err := strconv.Atoi(s)