coder · kylecarbs · May 17, 2023 · May 16, 2023 · May 16, 2023
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,7 +121,6 @@ func TestServer(t *testing.T) {
 		)
 
 		const superDuperLong = testutil.WaitSuperLong * 3
-
 		ctx := testutil.Context(t, superDuperLong)
 		clitest.Start(t, inv.WithContext(ctx))
 
@@ -1430,6 +1429,7 @@ func TestServer(t *testing.T) {
 					wantConfig.Options[i].Name,
 				)
 			}
+			w.Cancel()
 			w.RequireSuccess()
 		})
 	})
@@ -1460,8 +1460,8 @@ func TestServer(t *testing.T) {
 	})
 }
 
-//nolint:tparallel,paralleltest // This test spawns or connects to an existing PostgreSQL instance.
 func TestServer_Production(t *testing.T) {
+	t.Parallel()
 	if runtime.GOOS != "linux" || testing.Short() {
 		// Skip on non-Linux because it spawns a PostgreSQL instance.
 		t.SkipNow()
@@ -1471,7 +1471,8 @@ func TestServer_Production(t *testing.T) {
 	defer closeFunc()
 
 	// Postgres + race detector + CI = slow.
-	ctx := testutil.Context(t, testutil.WaitSuperLong*3)
+	ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitSuperLong*3)
+	defer cancelFunc()
 
 	inv, cfg := clitest.New(t,
 		"server",

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