fix: automatically cancel tests on cleanup

coder · kylecarbs · May 17, 2023 · May 16, 2023 · May 16, 2023 · May 16, 2023
commit 4f1916e470c6206b147b6dda2c68f604079251ec
diff --git a/cli/clitest/clitest.go b/cli/clitest/clitest.go
@@ -137,6 +137,7 @@ func Start(t *testing.T, inv *clibase.Invocation) {
 	// before ours.
 	waiter := StartWithWaiter(t, inv)
 	t.Cleanup(func() {
+		waiter.Cancel()
 		<-closeCh
 	})
 
@@ -163,11 +164,16 @@ func Run(t *testing.T, inv *clibase.Invocation) {
 type ErrorWaiter struct {
 	waitOnce    sync.Once
 	cachedError error
+	cancelFunc  context.CancelFunc
 
 	c <-chan error
 	t *testing.T
 }
 
+func (w *ErrorWaiter) Cancel() {
+	w.cancelFunc()
+}
+
 func (w *ErrorWaiter) Wait() error {
 	w.waitOnce.Do(func() {
 		var ok bool
@@ -241,5 +247,5 @@ func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
 		cleaningUp.Store(true)
 		<-doneCh
 	})
-	return &ErrorWaiter{c: errCh, t: t}
+	return &ErrorWaiter{c: errCh, t: t, cancelFunc: cancel}
 }
diff --git a/cli/gitssh_test.go b/cli/gitssh_test.go
@@ -58,9 +58,9 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*codersdk.Client, str
 
 	// start workspace agent
 	inv, root := clitest.New(t, "agent", "--agent-token", agentToken, "--agent-url", client.URL.String())
-	agentClient := client
+	agentClient := codersdk.New(client.URL)
+	agentClient.SetSessionToken(agentToken)
 	clitest.SetupConfig(t, agentClient, root)
-
 	clitest.Start(t, inv)
 
 	coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)

diff --git a/cli/server_test.go b/cli/server_test.go
@@ -121,9 +121,7 @@ func TestServer(t *testing.T) {
 		)
 
 		const superDuperLong = testutil.WaitSuperLong * 3
-
-		ctx, cancelFunc := context.WithTimeout(context.Background(), superDuperLong)
-		defer cancelFunc()
+		ctx := testutil.Context(t, superDuperLong)
 		clitest.Start(t, inv.WithContext(ctx))
 
 		//nolint:gocritic // Embedded postgres take a while to fire up.
@@ -1431,6 +1429,7 @@ func TestServer(t *testing.T) {
 					wantConfig.Options[i].Name,
 				)
 			}
+			w.Cancel()
 			w.RequireSuccess()
 		})
 	})

diff --git a/cli/vscodessh_test.go b/cli/vscodessh_test.go
@@ -68,6 +68,7 @@ func TestVSCodeSSH(t *testing.T) {
 		}
 		return len(entries) > 0
 	}, testutil.WaitLong, testutil.IntervalFast)
+	waiter.Cancel()
 
 	if err := waiter.Wait(); err != nil {
 		waiter.RequireIs(context.Canceled)

diff --git a/enterprise/coderd/workspaceproxy.go b/enterprise/coderd/workspaceproxy.go
@@ -31,7 +31,7 @@ import (
 // forceWorkspaceProxyHealthUpdate forces an update of the proxy health.
 // This is useful when a proxy is created or deleted. Errors will be logged.
 func (api *API) forceWorkspaceProxyHealthUpdate(ctx context.Context) {
-	if err := api.ProxyHealth.ForceUpdate(ctx); err != nil {
+	if err := api.ProxyHealth.ForceUpdate(ctx); err != nil && !xerrors.Is(err, context.Canceled) {
 		api.Logger.Error(ctx, "force proxy health update", slog.Error(err))
 	}
 }