decompose helper function to retain restart behavior

sreya · sreya · commit eb36e88d5d87 · 2023-11-02T00:37:35.000Z
diff --git a/cli/restart.go b/cli/restart.go
@@ -35,6 +35,16 @@ func (r *RootCmd) restart() *clibase.Cmd {
 				return err
 			}
 
+			startReq, err := workspaceStartRequest(inv, client, startWorkspaceArgs{
+				workspace:      workspace,
+				action:         WorkspaceRestart,
+				parameterFlags: parameterFlags,
+				alwaysPrompt:   true,
+			})
+			if err != nil {
+				return err
+			}
+
 			_, err = cliui.Prompt(inv, cliui.PromptOptions{
 				Text:      "Confirm restart workspace?",
 				IsConfirm: true,
@@ -55,20 +65,15 @@ func (r *RootCmd) restart() *clibase.Cmd {
 				return err
 			}
 
-			build, err = startWorkspace(inv, client, startWorkspaceArgs{
-				workspace:      workspace,
-				parameterFlags: parameterFlags,
-				action:         WorkspaceStart,
-				alwaysPrompt:   false,
-			})
+			build, err = client.CreateWorkspaceBuild(ctx, workspace.ID, startReq)
 			// It's possible for a workspace build to fail due to the template requiring starting
 			// workspaces with the active version.
 			if cerr, ok := codersdk.AsError(err); ok && cerr.StatusCode() == http.StatusUnauthorized {
 				build, err = startWorkspace(inv, client, startWorkspaceArgs{
 					workspace:      workspace,
 					parameterFlags: parameterFlags,
 					action:         WorkspaceUpdate,
-					alwaysPrompt:   false,
+					alwaysPrompt:   true,
 				})
 				if err != nil {
 					return xerrors.Errorf("start workspace with active template version: %w", err)
diff --git a/cli/start.go b/cli/start.go
@@ -75,29 +75,29 @@ type startWorkspaceArgs struct {
 	alwaysPrompt   bool
 }
 
-func startWorkspace(inv *clibase.Invocation, client *codersdk.Client, args startWorkspaceArgs) (codersdk.WorkspaceBuild, error) {
+func workspaceStartRequest(inv *clibase.Invocation, client *codersdk.Client, args startWorkspaceArgs) (codersdk.CreateWorkspaceBuildRequest, error) {
 	version := args.workspace.LatestBuild.TemplateVersionID
 	if args.workspace.AutomaticUpdates == codersdk.AutomaticUpdatesAlways || (args.action == WorkspaceUpdate || args.action == WorkspaceCreate) {
 		template, err := client.Template(inv.Context(), args.workspace.TemplateID)
 		if err != nil {
-			return codersdk.WorkspaceBuild{}, xerrors.Errorf("get template: %w", err)
+			return codersdk.CreateWorkspaceBuildRequest{}, xerrors.Errorf("get template: %w", err)
 		}
 		version = template.ActiveVersionID
 	}
 
 	lastBuildParameters, err := client.WorkspaceBuildParameters(inv.Context(), args.workspace.LatestBuild.ID)
 	if err != nil {
-		return codersdk.WorkspaceBuild{}, err
+		return codersdk.CreateWorkspaceBuildRequest{}, err
 	}
 
 	buildOptions, err := asWorkspaceBuildParameters(args.parameterFlags.buildOptions)
 	if err != nil {
-		return codersdk.WorkspaceBuild{}, xerrors.Errorf("unable to parse build options: %w", err)
+		return codersdk.CreateWorkspaceBuildRequest{}, xerrors.Errorf("unable to parse build options: %w", err)
 	}
 
 	cliRichParameters, err := asWorkspaceBuildParameters(args.parameterFlags.richParameters)
 	if err != nil {
-		return codersdk.WorkspaceBuild{}, xerrors.Errorf("unable to parse build options: %w", err)
+		return codersdk.CreateWorkspaceBuildRequest{}, xerrors.Errorf("unable to parse build options: %w", err)
 	}
 
 	buildParameters, err := prepWorkspaceBuild(inv, client, prepWorkspaceBuildArgs{
@@ -113,13 +113,20 @@ func startWorkspace(inv *clibase.Invocation, client *codersdk.Client, args start
 		RichParameterFile:    args.parameterFlags.richParameterFile,
 	})
 	if err != nil {
-		return codersdk.WorkspaceBuild{}, err
+		return codersdk.CreateWorkspaceBuildRequest{}, err
 	}
 
-	req := codersdk.CreateWorkspaceBuildRequest{
+	return codersdk.CreateWorkspaceBuildRequest{
 		Transition:          codersdk.WorkspaceTransitionStart,
 		RichParameterValues: buildParameters,
 		TemplateVersionID:   version,
+	}, nil
+}
+
+func startWorkspace(inv *clibase.Invocation, client *codersdk.Client, args startWorkspaceArgs) (codersdk.WorkspaceBuild, error) {
+	req, err := workspaceStartRequest(inv, client, args)
+	if err != nil {
+		return codersdk.WorkspaceBuild{}, err
 	}
 
 	build, err := client.CreateWorkspaceBuild(inv.Context(), args.workspace.ID, req)
