fix: automatically cancel tests on cleanup

kylecarbs · kylecarbs · commit b62b27abde26 · 2023-05-16T17:00:16.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
@@ -1431,6 +1431,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)

Original file line number	Diff line number	Diff line change
`@@ -1431,6 +1431,7 @@ func TestServer(t *testing.T) {`
`1431`	`1431`	`wantConfig.Options[i].Name,`
`1432`	`1432`	`)`
`1433`	`1433`	`}`
	`1434`	`+ w.Cancel()`
`1434`	`1435`	`w.RequireSuccess()`
`1435`	`1436`	`})`
`1436`	`1437`	`})`
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)`