Code Optimization

choice/choice.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"strings"
@@ -29,17 +30,13 @@ func Menu[E any](choices []E, showQuit bool) string {
 	if len(choices) == 0 {
 		return ""
 	}
-	var digit int
-	for n := len(choices); n != 0; digit++ {
-		n /= 10
-	}
-	option := fmt.Sprintf("%%%dd", digit)
+	digit := len(strconv.Itoa(len(choices)))
 	var b strings.Builder
 	for i, choice := range choices {
-		fmt.Fprintf(&b, "%s. %s\n", fmt.Sprintf(option, i+1), choiceStr(choice))
+		fmt.Fprintf(&b, "%*d. %s\n", digit, i+1, choiceStr(choice))
 	}
 	if showQuit {
-		fmt.Fprintf(&b, "%s. Quit\n", fmt.Sprintf(fmt.Sprintf("%%%ds", digit), "q"))
+		fmt.Fprintf(&b, "%*d. Quit\n", digit, 0)
 	}
 	return b.String()
 }
@@ -51,13 +48,8 @@ var _ error = choiceError("")
 
 type choiceError string
 
-func (err choiceError) Error() string {
-	return "bad choice: " + string(err)
-}
-
-func (choiceError) Unwrap() error {
-	return ErrBadChoice
-}
+func (err choiceError) Error() string { return "bad choice: " + string(err) }
+func (choiceError) Unwrap() error     { return ErrBadChoice }
 
 func choose[E any](choice string, choices []E) (res E, err error) {
 	n, err := strconv.Atoi(choice)
@@ -79,6 +71,10 @@ func Choose[E any](choices []E) (choice bool, res E, err error) {
 
 // ChooseWithDefault function allows the user to make a choice from the given options with an optional default value.
 func ChooseWithDefault[E any](choices []E, def int) (choice bool, res E, err error) {
+	return chooseWithDefault(os.Stdin, choices, def)
+}
+
+func chooseWithDefault[E any](r io.Reader, choices []E, def int) (choice bool, res E, err error) {
 	if n := len(choices); n == 0 {
 		err = errors.New("no choices")
 		return
@@ -92,26 +88,31 @@ func ChooseWithDefault[E any](choices []E, def int) (choice bool, res E, err err
 	} else {
 		prompt = "Please choose: "
 	}
-	scanner := bufio.NewScanner(os.Stdin)
-	var b []byte
-	if def <= 0 {
-		for len(b) == 0 {
-			fmt.Print(prompt)
-			scanner.Scan()
-			b = bytes.TrimSpace(scanner.Bytes())
-		}
-	} else {
-		fmt.Print(prompt)
-		scanner.Scan()
-		b = bytes.TrimSpace(scanner.Bytes())
-		if len(b) == 0 {
-			return true, choices[def-1], nil
-		}
+	b, err := readLine(bufio.NewScanner(r), prompt, def <= 0)
+	if err != nil {
+		return
+	}
+	if len(b) == 0 && def > 0 {
+		return true, choices[def-1], nil
 	}
-	if bytes.EqualFold(b, []byte("q")) {
+	if bytes.EqualFold(b, []byte("0")) || bytes.EqualFold(b, []byte("q")) {
 		return
 	}
 	choice = true
 	res, err = choose(string(b), choices)
 	return
 }
+
+func readLine(scanner *bufio.Scanner, prompt string, required bool) ([]byte, error) {
+	for {
+		fmt.Print(prompt)
+		if !scanner.Scan() {
+			return nil, scanner.Err()
+		}
+		b := bytes.TrimSpace(scanner.Bytes())
+		if required && len(b) == 0 {
+			continue
+		}
+		return b, nil
+	}
+}

choice/choice_test.go

@@ -30,7 +30,7 @@ func TestMenu(t *testing.T) {
  8. hh
  9. ii
 10. jj
- q. Quit
+ 0. Quit
 `
 	if s := Menu(choices, true); s != expect {
 		t.Errorf("expected %q; got %q", expect, s)
@@ -54,3 +54,104 @@ func TestError(t *testing.T) {
 		t.Error("expected err is ErrBadChoice; got not")
 	}
 }
+
+func TestChooseWithDefault(t *testing.T) {
+	tests := []struct {
+		name       string
+		input      string
+		choices    []string
+		def        int
+		wantChoice bool
+		wantRes    string
+		wantErr    bool
+	}{
+		{
+			name:       "valid choice",
+			input:      "2\n",
+			choices:    []string{"a", "b", "c"},
+			def:        0,
+			wantChoice: true,
+			wantRes:    "b",
+			wantErr:    false,
+		},
+		{
+			name:       "default used when empty input",
+			input:      "\n",
+			choices:    []string{"a", "b", "c"},
+			def:        2,
+			wantChoice: true,
+			wantRes:    "b",
+			wantErr:    false,
+		},
+		{
+			name:       "quit with 0",
+			input:      "0\n",
+			choices:    []string{"a", "b"},
+			def:        0,
+			wantChoice: false,
+			wantRes:    "",
+			wantErr:    false,
+		},
+		{
+			name:       "quit with q",
+			input:      "q\n",
+			choices:    []string{"a", "b"},
+			def:        0,
+			wantChoice: false,
+			wantRes:    "",
+			wantErr:    false,
+		},
+		{
+			name:       "invalid input",
+			input:      "x\n",
+			choices:    []string{"a", "b"},
+			def:        0,
+			wantChoice: true,
+			wantRes:    "",
+			wantErr:    true,
+		},
+		{
+			name:       "out of range",
+			input:      "5\n",
+			choices:    []string{"a", "b"},
+			def:        0,
+			wantChoice: true,
+			wantRes:    "",
+			wantErr:    true,
+		},
+		{
+			name:       "no choices",
+			input:      "1\n",
+			choices:    []string{},
+			def:        0,
+			wantChoice: false,
+			wantRes:    "",
+			wantErr:    true,
+		},
+		{
+			name:       "invalid default",
+			input:      "1\n",
+			choices:    []string{"a"},
+			def:        5,
+			wantChoice: false,
+			wantRes:    "",
+			wantErr:    true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := strings.NewReader(tt.input)
+			choice, res, err := chooseWithDefault(r, tt.choices, tt.def)
+			if choice != tt.wantChoice {
+				t.Errorf("got choice=%v, want %v", choice, tt.wantChoice)
+			}
+			if res != tt.wantRes {
+				t.Errorf("got res=%q, want %q", res, tt.wantRes)
+			}
+			if (err != nil) != tt.wantErr {
+				t.Errorf("got err=%v, wantErr=%v", err, tt.wantErr)
+			}
+		})
+	}
+}