mitchellh · mitchellh · Mar 26, 2020 · Dec 25, 2017
diff --git a/ui_mock.go b/ui_mock.go
@@ -1,9 +1,11 @@
 package cli
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
+	"strings"
 	"sync"
 )
 
@@ -35,9 +37,12 @@ func (u *MockUi) Ask(query string) (string, error) {
 
 	var result string
 	fmt.Fprint(u.OutputWriter, query)
-	if _, err := fmt.Fscanln(u.InputReader, &result); err != nil {
+	r := bufio.NewReader(u.InputReader)
+	line, err := r.ReadString('\n')
+	if err != nil {
 		return "", err
 	}
+	result = strings.TrimRight(line, "\r\n")
 
 	return result, nil
 }

diff --git a/ui_mock_test.go b/ui_mock_test.go
@@ -1,9 +1,44 @@
 package cli
 
 import (
+	"io"
 	"testing"
 )
 
 func TestMockUi_implements(t *testing.T) {
 	var _ Ui = new(MockUi)
 }
+
+func TestMockUi_Ask(t *testing.T) {
+	tests := []struct {
+		name           string
+		query, input   string
+		expectedResult string
+	}{
+		{"EmptyString", "Middle Name?", "\n", ""},
+		{"NonEmptyString", "Name?", "foo bar\nbaz\n", "foo bar"},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			in_r, in_w := io.Pipe()
+			defer in_r.Close()
+			defer in_w.Close()
+
+			ui := &MockUi{
+				InputReader: in_r,
+			}
+
+			go in_w.Write([]byte(tc.input))
+
+			result, err := ui.Ask(tc.query)
+			if err != nil {
+				t.Fatalf("err: %s", err)
+			}
+
+			if result != tc.expectedResult {
+				t.Fatalf("bad: %#v", result)
+			}
+		})
+	}
+}
diff --git a/ui_test.go b/ui_test.go
@@ -11,29 +11,42 @@ func TestBasicUi_implements(t *testing.T) {
 }
 
 func TestBasicUi_Ask(t *testing.T) {
-	in_r, in_w := io.Pipe()
-	defer in_r.Close()
-	defer in_w.Close()
-
-	writer := new(bytes.Buffer)
-	ui := &BasicUi{
-		Reader: in_r,
-		Writer: writer,
-	}
-
-	go in_w.Write([]byte("foo bar\nbaz\n"))
-
-	result, err := ui.Ask("Name?")
-	if err != nil {
-		t.Fatalf("err: %s", err)
-	}
-
-	if writer.String() != "Name? " {
-		t.Fatalf("bad: %#v", writer.String())
-	}
-
-	if result != "foo bar" {
-		t.Fatalf("bad: %#v", result)
+	tests := []struct {
+		name                          string
+		query, input                  string
+		expectedQuery, expectedResult string
+	}{
+		{"EmptyString", "Middle Name?", "\n", "Middle Name? ", ""},
+		{"NonEmptyString", "Name?", "foo bar\nbaz\n", "Name? ", "foo bar"},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			in_r, in_w := io.Pipe()
+			defer in_r.Close()
+			defer in_w.Close()
+
+			writer := new(bytes.Buffer)
+			ui := &BasicUi{
+				Reader: in_r,
+				Writer: writer,
+			}
+
+			go in_w.Write([]byte(tc.input))
+
+			result, err := ui.Ask(tc.query)
+			if err != nil {
+				t.Fatalf("err: %s", err)
+			}
+
+			if writer.String() != tc.expectedQuery {
+				t.Fatalf("bad: %#v", writer.String())
+			}
+
+			if result != tc.expectedResult {
+				t.Fatalf("bad: %#v", result)
+			}
+		})
 	}
 }