test(cli): prevent flake due to outdated build in TestSSH (coder#12760)

mafredri · web-flow · commit f418ece9ae23 · 2024-03-26T10:46:58.000Z
Fixes coder#12752
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -117,14 +117,26 @@ func TestSSH(t *testing.T) {
 		clitest.SetupConfig(t, client, root)
 		pty := ptytest.New(t).Attach(inv)
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
 		defer cancel()
 
 		cmdDone := tGo(t, func() {
 			err := inv.WithContext(ctx).Run()
 			assert.NoError(t, err)
 		})
 
+		// Delay until workspace is starting, otherwise the agent may be
+		// booted due to outdated build.
+		var err error
+		for {
+			workspace, err = client.Workspace(ctx, workspace.ID)
+			require.NoError(t, err)
+			if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart {
+				break
+			}
+			time.Sleep(testutil.IntervalFast)
+		}
+
 		// When the agent connects, the workspace was started, and we should
 		// have access to the shell.
 		_ = agenttest.New(t, client.URL, authToken)
@@ -365,7 +377,7 @@ func TestSSH(t *testing.T) {
 		workspaceBuild := coderdtest.CreateWorkspaceBuild(t, client, workspace, database.WorkspaceTransitionStop)
 		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspaceBuild.ID)
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitSuperLong)
 		defer cancel()
 
 		clientStdinR, clientStdinW := io.Pipe()
@@ -461,6 +473,18 @@ func TestSSH(t *testing.T) {
 			assert.NoError(t, err)
 		})
 
+		// Delay until workspace is starting, otherwise the agent may be
+		// booted due to outdated build.
+		var err error
+		for {
+			workspace, err = client.Workspace(ctx, workspace.ID)
+			require.NoError(t, err)
+			if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart {
+				break
+			}
+			time.Sleep(testutil.IntervalFast)
+		}
+
 		// When the agent connects, the workspace was started, and we should
 		// have access to the shell.
 		_ = agenttest.New(t, client.URL, authToken)
