fix: automatically cancel tests on cleanup

kylecarbs · kylecarbs · commit 4f1916e470c6 · 2023-05-16T17:34:45.000Z
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))
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -121,9 +121,7 @@ func TestServer(t *testing.T) {`
`121`	`121`	`)`
`122`	`122`
`123`	`123`	`const superDuperLong = testutil.WaitSuperLong * 3`
`124`		`-`
`125`		`- ctx, cancelFunc := context.WithTimeout(context.Background(), superDuperLong)`
`126`		`- defer cancelFunc()`
	`124`	`+ ctx := testutil.Context(t, superDuperLong)`
`127`	`125`	`clitest.Start(t, inv.WithContext(ctx))`
`128`	`126`
`129`	`127`	`//nolint:gocritic // Embedded postgres take a while to fire up.`
`@@ -1431,6 +1429,7 @@ func TestServer(t *testing.T) {`
`1431`	`1429`	`wantConfig.Options[i].Name,`
`1432`	`1430`	`)`
`1433`	`1431`	`}`
	`1432`	`+ w.Cancel()`
`1434`	`1433`	`w.RequireSuccess()`
`1435`	`1434`	`})`
`1436`	`1435`	`})`
Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,7 @@ func TestVSCodeSSH(t *testing.T) {`
`68`	`68`	`}`
`69`	`69`	`return len(entries) > 0`
`70`	`70`	`}, testutil.WaitLong, testutil.IntervalFast)`
	`71`	`+ waiter.Cancel()`
`71`	`72`
`72`	`73`	`if err := waiter.Wait(); err != nil {`
`73`	`74`	`waiter.RequireIs(context.Canceled)`
Original file line number	Diff line number	Diff line change
`@@ -31,7 +31,7 @@ import (`
`31`	`31`	`// forceWorkspaceProxyHealthUpdate forces an update of the proxy health.`
`32`	`32`	`// This is useful when a proxy is created or deleted. Errors will be logged.`
`33`	`33`	`func (api *API) forceWorkspaceProxyHealthUpdate(ctx context.Context) {`
`34`		`- if err := api.ProxyHealth.ForceUpdate(ctx); err != nil {`
	`34`	`+ if err := api.ProxyHealth.ForceUpdate(ctx); err != nil && !xerrors.Is(err, context.Canceled) {`
`35`	`35`	`api.Logger.Error(ctx, "force proxy health update", slog.Error(err))`
`36`	`36`	`}`
`37`	`37`	`}`