Wow, I made the tests pass

kylecarbs · kylecarbs · commit 6fa941f958ab · 2022-10-14T16:16:08.000Z
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -465,7 +465,7 @@ func TestAgent(t *testing.T) {
 
 				conn, _ := setupAgent(t, codersdk.WorkspaceAgentMetadata{}, 0)
 				require.Eventually(t, func() bool {
-					_, err := conn.Ping()
+					_, err := conn.Ping(context.Background())
 					return err == nil
 				}, testutil.WaitMedium, testutil.IntervalFast)
 				conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String())
diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -69,7 +69,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.NoError(t, err)
 		defer dialer.Close()
 		require.Eventually(t, func() bool {
-			_, err := dialer.Ping()
+			_, err := dialer.Ping(ctx)
 			return err == nil
 		}, testutil.WaitMedium, testutil.IntervalFast)
 		cancelFunc()
@@ -130,7 +130,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.NoError(t, err)
 		defer dialer.Close()
 		require.Eventually(t, func() bool {
-			_, err := dialer.Ping()
+			_, err := dialer.Ping(ctx)
 			return err == nil
 		}, testutil.WaitMedium, testutil.IntervalFast)
 		cancelFunc()
@@ -191,7 +191,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.NoError(t, err)
 		defer dialer.Close()
 		require.Eventually(t, func() bool {
-			_, err := dialer.Ping()
+			_, err := dialer.Ping(ctx)
 			return err == nil
 		}, testutil.WaitMedium, testutil.IntervalFast)
 		cancelFunc()
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -155,7 +155,7 @@ func portForward() *cobra.Command {
 				case <-ticker.C:
 				}
 
-				_, err = conn.Ping()
+				_, err = conn.Ping(ctx)
 				if err != nil {
 					continue
 				}
diff --git a/cli/speedtest.go b/cli/speedtest.go
@@ -70,7 +70,7 @@ func speedtest() *cobra.Command {
 					return ctx.Err()
 				case <-ticker.C:
 				}
-				dur, err := conn.Ping()
+				dur, err := conn.Ping(ctx)
 				if err != nil {
 					continue
 				}
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -129,7 +129,7 @@ func TestWorkspaceAgentListen(t *testing.T) {
 			_ = conn.Close()
 		}()
 		require.Eventually(t, func() bool {
-			_, err := conn.Ping()
+			_, err := conn.Ping(ctx)
 			return err == nil
 		}, testutil.WaitLong, testutil.IntervalFast)
 	})
diff --git a/codersdk/agentconn.go b/codersdk/agentconn.go
@@ -132,7 +132,7 @@ type AgentConn struct {
 	CloseFunc func()
 }
 
-func (c *AgentConn) Ping() (time.Duration, error) {
+func (c *AgentConn) Ping(ctx context.Context) (time.Duration, error) {
 	errCh := make(chan error, 1)
 	durCh := make(chan time.Duration, 1)
 	c.Conn.Ping(TailnetIP, tailcfg.PingDisco, func(pr *ipnstate.PingResult) {
@@ -145,6 +145,8 @@ func (c *AgentConn) Ping() (time.Duration, error) {
 	select {
 	case err := <-errCh:
 		return 0, err
+	case <-ctx.Done():
+		return 0, ctx.Err()
 	case dur := <-durCh:
 		return dur, nil
 	}
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -3,7 +3,6 @@ package coderd
 import (
 	"context"
 	"crypto/ed25519"
-	"fmt"
 	"net/http"
 	"sync"
 	"time"
@@ -247,7 +246,6 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 				coordinator = haCoordinator
 			}
 
-			fmt.Printf("HA enabled\n")
 			api.replicaManager.SetCallback(func() {
 				addresses := make([]string, 0)
 				for _, replica := range api.replicaManager.Regional() {
diff --git a/enterprise/coderd/replicas_test.go b/enterprise/coderd/replicas_test.go
@@ -3,6 +3,7 @@ package coderd_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -69,7 +70,9 @@ func TestReplicas(t *testing.T) {
 		})
 		require.NoError(t, err)
 		require.Eventually(t, func() bool {
-			_, err = conn.Ping()
+			ctx, cancelFunc := context.WithTimeout(context.Background(), 3*time.Second)
+			defer cancelFunc()
+			_, err = conn.Ping(ctx)
 			return err == nil
 		}, testutil.WaitLong, testutil.IntervalFast)
 		_ = conn.Close()

Original file line number	Diff line number	Diff line change
`@@ -155,7 +155,7 @@ func portForward() *cobra.Command {`
`155`	`155`	`case <-ticker.C:`
`156`	`156`	`}`
`157`	`157`
`158`		`- _, err = conn.Ping()`
	`158`	`+ _, err = conn.Ping(ctx)`
`159`	`159`	`if err != nil {`
`160`	`160`	`continue`
`161`	`161`	`}`
Original file line number	Diff line number	Diff line change
`@@ -70,7 +70,7 @@ func speedtest() *cobra.Command {`
`70`	`70`	`return ctx.Err()`
`71`	`71`	`case <-ticker.C:`
`72`	`72`	`}`
`73`		`- dur, err := conn.Ping()`
	`73`	`+ dur, err := conn.Ping(ctx)`
`74`	`74`	`if err != nil {`
`75`	`75`	`continue`
`76`	`76`	`}`