Improve coverage, revert zero duration change

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)
 			}

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())

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)