fix: Improve output, detect if stdout is tty

mafredri · mafredri · commit c64686753225 · 2022-05-31T11:34:44.000+03:00
diff --git a/cli/configssh.go b/cli/configssh.go
@@ -103,6 +103,15 @@ func configSSH() *cobra.Command {
 			// TODO(mafredri): Check coderConfigFile for previous options
 			// coderConfigFile.
 
+			out := cmd.OutOrStdout()
+			if showDiff {
+				out = cmd.OutOrStderr()
+			}
+			binaryFile, err := currentBinPath(out)
+			if err != nil {
+				return err
+			}
+
 			// Only allow not-exist errors to avoid trashing
 			// the users SSH config.
 			configRaw, err := os.ReadFile(sshConfigFile)
@@ -136,11 +145,6 @@ func configSSH() *cobra.Command {
 				configModified = append(configModified, []byte(sep+sshConfigIncludeStatement+"\n")...)
 			}
 
-			binaryFile, err := currentBinPath(cmd)
-			if err != nil {
-				return err
-			}
-
 			root := createConfig(cmd)
 			var errGroup errgroup.Group
 			type workspaceConfig struct {
@@ -240,13 +244,13 @@ func configSSH() *cobra.Command {
 			if showDiff {
 				if len(changes) > 0 {
 					// Write to stderr to avoid dirtying the diff output.
-					_, _ = fmt.Fprint(cmd.ErrOrStderr(), "Changes:\n\n")
+					_, _ = fmt.Fprint(out, "Changes:\n\n")
 					for _, change := range changes {
-						_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "* %s\n", change)
+						_, _ = fmt.Fprintf(out, "* %s\n", change)
 					}
 				}
 
-				color := isTTY(cmd)
+				color := isTTYOut(cmd)
 				for _, diffFn := range []func() ([]byte, error){
 					func() ([]byte, error) { return diffBytes(sshConfigFile, configRaw, configModified, color) },
 					func() ([]byte, error) { return diffBytes(coderConfigFile, coderConfigRaw, buf.Bytes(), color) },
@@ -257,6 +261,7 @@ func configSSH() *cobra.Command {
 					}
 					if len(diff) > 0 {
 						filesDiffer = true
+						// Always write to stdout.
 						_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\n%s", diff)
 					}
 				}
@@ -270,9 +275,9 @@ func configSSH() *cobra.Command {
 					IsConfirm: true,
 				})
 				if err != nil {
-					return err
+					return nil
 				}
-				_, _ = fmt.Fprint(cmd.OutOrStdout(), "\n")
+				_, _ = fmt.Fprint(out, "\n")
 
 				if !bytes.Equal(configRaw, configModified) {
 					err = writeWithTempFileAndMove(sshConfigFile, bytes.NewReader(configRaw))
@@ -289,10 +294,10 @@ func configSSH() *cobra.Command {
 			}
 
 			if len(workspaces) > 0 {
-				_, _ = fmt.Fprintln(cmd.OutOrStdout(), "You should now be able to ssh into your workspace")
-				_, _ = fmt.Fprintf(cmd.OutOrStdout(), "For example, try running:\n\n\t$ ssh coder.%s\n\n", workspaces[0].Name)
+				_, _ = fmt.Fprintln(out, "You should now be able to ssh into your workspace.")
+				_, _ = fmt.Fprintf(out, "For example, try running:\n\n\t$ ssh coder.%s\n\n", workspaces[0].Name)
 			} else {
-				_, _ = fmt.Fprint(cmd.OutOrStdout(), "You don't have any workspaces yet, try creating one with:\n\n\t$ coder create <workspace>\n\n")
+				_, _ = fmt.Fprint(out, "You don't have any workspaces yet, try creating one with:\n\n\t$ coder create <workspace>\n\n")
 			}
 			return nil
 		},
@@ -349,7 +354,7 @@ func writeWithTempFileAndMove(path string, r io.Reader) (err error) {
 
 // currentBinPath returns the path to the coder binary suitable for use in ssh
 // ProxyCommand.
-func currentBinPath(cmd *cobra.Command) (string, error) {
+func currentBinPath(w io.Writer) (string, error) {
 	exePath, err := os.Executable()
 	if err != nil {
 		return "", xerrors.Errorf("get executable path: %w", err)
@@ -365,24 +370,26 @@ func currentBinPath(cmd *cobra.Command) (string, error) {
 	// correctly. Check if the current executable is in $PATH, and warn the user
 	// if it isn't.
 	if err != nil && runtime.GOOS == "windows" {
-		cliui.Warn(cmd.OutOrStdout(),
+		cliui.Warn(w,
 			"The current executable is not in $PATH.",
 			"This may lead to problems connecting to your workspace via SSH.",
 			fmt.Sprintf("Please move %q to a location in your $PATH (such as System32) and run `%s config-ssh` again.", binName, binName),
 		)
+		_, _ = fmt.Fprint(w, "\n")
 		// Return the exePath so SSH at least works outside of Msys2.
 		return exePath, nil
 	}
 
 	// Warn the user if the current executable is not the same as the one in
 	// $PATH.
 	if filepath.Clean(pathPath) != filepath.Clean(exePath) {
-		cliui.Warn(cmd.OutOrStdout(),
+		cliui.Warn(w,
 			"The current executable path does not match the executable path found in $PATH.",
 			"This may cause issues connecting to your workspace via SSH.",
 			fmt.Sprintf("\tCurrent executable path: %q", exePath),
 			fmt.Sprintf("\tExecutable path in $PATH: %q", pathPath),
 		)
+		_, _ = fmt.Fprint(w, "\n")
 	}
 
 	return binName, nil
diff --git a/cli/root.go b/cli/root.go
@@ -203,6 +203,24 @@ func isTTY(cmd *cobra.Command) bool {
 	return isatty.IsTerminal(file.Fd())
 }
 
+// isTTYOut returns whether the passed reader is a TTY or not.
+// This accepts a reader to work with Cobra's "OutOrStdout"
+// function for simple testing.
+func isTTYOut(cmd *cobra.Command) bool {
+	// If the `--force-tty` command is available, and set,
+	// assume we're in a tty. This is primarily for cases on Windows
+	// where we may not be able to reliably detect this automatically (ie, tests)
+	forceTty, err := cmd.Flags().GetBool(varForceTty)
+	if forceTty && err == nil {
+		return true
+	}
+	file, ok := cmd.OutOrStdout().(*os.File)
+	if !ok {
+		return false
+	}
+	return isatty.IsTerminal(file.Fd())
+}
+
 func usageTemplate() string {
 	// usageHeader is defined in init().
 	return `{{usageHeader "Usage:"}}
