fix: tests

Author: fiam

commands/run.go

@@ -167,9 +167,8 @@ func newRunCmd() *cobra.Command {
 				case strings.HasPrefix(question, "/"):
 					printHelp(helpUnknownCommand)
 
-				case strings.HasPrefix(question, `"""`):
-					initialText := strings.TrimPrefix(question, `"""`)
-					restOfText, err := readMultilineString(cmd.Context(), initialText)
+				case strings.HasPrefix(question, `"""`) || strings.HasPrefix(question, `'''`):
+					restOfText, err := readMultilineString(cmd.Context(), os.Stdin, question)
 					if err != nil {
 						return err
 					}
@@ -323,54 +322,72 @@ func copyToClipboard(command string) error {
 	return nil
 }
 
-// readMultilineString reads a multiline string from stdin, starting with the initial text if provided.
-// The initialText parameter is optional and can be used to provide a starting point for the input.
-// It reads from stdin until it encounters a closing triple quote (""").
-func readMultilineString(ctx context.Context, initialText string) (string, error) {
+// readMultilineString reads a multiline string from r, starting with the initial text if provided.
+// The text must start either with triple quotes (""") or single quotes (”'). readMultilineString
+// will scan the input until a matching closing quote is found, or return an error if the input
+// is not properly closed. It returns the string content without the surrounding quotes but
+// preserving the original newlines and indentation.
+func readMultilineString(ctx context.Context, r io.Reader, initialText string) (string, error) {
 	var question string
 
 	// Start with the initial text if provided
 	if initialText != "" {
-		question = initialText + "\n"
+		tr := strings.NewReader(initialText)
+		mr := io.MultiReader(tr, r)
+		r = mr
+	}
+
+	br := bufio.NewReader(r)
+
+	// Read the first 3 bytes
+	pd := make([]byte, 3)
+	if _, err := io.ReadFull(br, pd); err != nil {
+		return "", err
+	}
+	prefix := string(pd)
+	if prefix != `"""` && prefix != `'''` {
+		return "", fmt.Errorf("invalid multiline string prefix: %q", prefix)
 	}
 
 	for {
 		fmt.Print("... ")
 
-		line, err := readLine(ctx)
+		line, err := readLine(ctx, br)
 		if err != nil {
+			if errors.Is(err, io.EOF) {
+				return "", fmt.Errorf("unclosed multiline input: %w", err)
+			}
 			return "", err
 		}
 
-		question += line + "\n"
-		if strings.TrimSpace(line) == `"""` || strings.HasSuffix(strings.TrimSpace(line), `"""`) {
+		question += line
+		if strings.TrimSpace(line) == prefix || strings.HasSuffix(strings.TrimSpace(line), prefix) {
 			break
 		}
 	}
 
-	// Find and remove the closing triple quotes
+	// Find and remove the closing triple terminator
 	content := question
-	if idx := strings.LastIndex(content, `"""`); idx >= 0 {
+	if idx := strings.LastIndex(content, prefix); idx >= 0 {
 		before := content[:idx]
-		after := content[idx+3:]
+		after := content[idx+len(prefix):]
 		content = before + after
 	}
 
 	return strings.TrimRight(content, "\n"), nil
 }
 
-// readLine reads a single line from stdin.
-func readLine(ctx context.Context) (string, error) {
+// readLine reads a single line from r.
+func readLine(ctx context.Context, r *bufio.Reader) (string, error) {
 	lines := make(chan string)
 	errs := make(chan error)
 
 	go func() {
 		defer close(lines)
 		defer close(errs)
 
-		reader := bufio.NewReader(os.Stdin)
-		line, err := reader.ReadString('\n')
-		if err != nil {
+		line, err := r.ReadString('\n')
+		if err != nil && line == "" {
 			errs <- err
 		} else {
 			lines <- line

commands/run_test.go

@@ -2,10 +2,13 @@ package commands
 
 import (
 	"bufio"
+	"errors"
+	"io"
 	"strings"
 	"testing"
 
-	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestReadMultilineInput(t *testing.T) {
@@ -15,40 +18,34 @@ func TestReadMultilineInput(t *testing.T) {
 		expected string
 		wantErr  bool
 	}{
-		{
-			name:     "single line input",
-			input:    "hello world",
-			expected: "hello world",
-			wantErr:  false,
-		},
 		{
 			name:     "single line with triple quotes",
 			input:    `"""hello world"""`,
-			expected: `"""hello world"""`,
+			expected: `hello world`,
 			wantErr:  false,
 		},
 		{
 			name: "multiline input with double quotes",
 			input: `"""tell
-me
-a
-joke"""`,
-			expected: `"""tell
-me
-a
-joke"""`,
+			me
+			a
+			joke"""`,
+			expected: `tell
+			me
+			a
+			joke`,
 			wantErr: false,
 		},
 		{
 			name: "multiline input with single quotes",
 			input: `'''tell
-me
-a
-joke'''`,
-			expected: `'''tell
-me
-a
-joke'''`,
+			me
+			a
+			joke'''`,
+			expected: `tell
+			me
+			a
+			joke`,
 			wantErr: false,
 		},
 		{
@@ -61,27 +58,124 @@ joke'''`,
 			name: "multiline with empty lines",
 			input: `"""first line
 
-third line"""`,
-			expected: `"""first line
+			third line"""`,
+			expected: `first line
 
-third line"""`,
+			third line`,
+			wantErr: false,
+		},
+		{
+			name: "multiline with spaces and closing quotes on new line",
+			input: `"""first line
+second line
+third line
+"""`,
+			expected: `first line
+second line
+third line`, // this will intentionally trim the last newline
+			wantErr: false,
+		},
+		{
+			name: "multiline with closing quotes and trailing spaces",
+			input: `"""first line
+second line
+third line   """`,
+			expected: `first line
+second line
+third line   `,
+			wantErr: false,
+		},
+		{
+			name:     "single quotes with spaces",
+			input:    `'''foo bar'''`,
+			expected: `foo bar`,
+			wantErr:  false,
+		},
+		{
+			name:     "triple quotes only",
+			input:    `""""""`,
+			expected: "",
+			wantErr:  false,
+		},
+		{
+			name:     "single quotes only",
+			input:    `''''''`,
+			expected: "",
+			wantErr:  false,
+		},
+		{
+			name:     "closing quotes in middle of line",
+			input:    `"""foo"""bar"""`,
+			expected: `foo"""bar`,
+			wantErr:  false,
+		},
+		{
+			name: "no closing quotes",
+			input: `"""foo
+bar
+baz`,
+			expected: "",
+			wantErr:  true,
+		},
+		{
+			name:     "invalid prefix",
+			input:    `"foo bar"`,
+			expected: "",
+			wantErr:  true,
+		},
+		{
+			name:     "prefix but no content",
+			input:    `"""`,
+			expected: "",
+			wantErr:  true,
+		},
+		{
+			name: "prefix and newline only",
+			input: `"""
+"""`,
+			expected: "",
+			wantErr:  false,
+		},
+		{
+			name: "multiline with only whitespace",
+			input: `"""   
+   
+   """`,
+			expected: `   
+   
+   `,
 			wantErr: false,
 		},
 	}
 
+	// Strings can be read either as the first line followed by the rest of the text,
+	// or as a single block of text. Make sure we test both scenarios.
+
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			// Create a mock command for testing
-			cmd := &cobra.Command{}
+			r := bufio.NewReader(strings.NewReader(tt.input))
+			result, err := readMultilineString(t.Context(), r, "")
 
-			// Create a scanner from the test input
-			scanner := bufio.NewScanner(strings.NewReader(tt.input))
+			if (err != nil) != tt.wantErr {
+				t.Errorf("readMultilineInput() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
 
-			// Capture output to avoid printing during tests
-			var output strings.Builder
-			cmd.SetOut(&output)
+			if result != tt.expected {
+				t.Errorf("readMultilineInput() = %q, want %q", result, tt.expected)
+			}
+		})
 
-			result, err := readMultilineInput(cmd, scanner)
+		t.Run(tt.name+"_chunked", func(t *testing.T) {
+			r := bufio.NewReader(strings.NewReader(tt.input))
+			firstLine, err := r.ReadString('\n') // Simulate reading the first line
+			if errors.Is(err, io.EOF) {
+				// Some test cases are single line, EOF is ok here
+				firstLine = tt.input
+			} else {
+				require.NoError(t, err)
+			}
+			result, err := readMultilineString(t.Context(), r, firstLine)
 
 			if (err != nil) != tt.wantErr {
 				t.Errorf("readMultilineInput() error = %v, wantErr %v", err, tt.wantErr)
@@ -98,18 +192,12 @@ third line"""`,
 func TestReadMultilineInputUnclosed(t *testing.T) {
 	// Test unclosed multiline input (should return error)
 	input := `"""unclosed multiline`
-	cmd := &cobra.Command{}
-	var output strings.Builder
-	cmd.SetOut(&output)
-
-	scanner := bufio.NewScanner(strings.NewReader(input))
-
-	_, err := readMultilineInput(cmd, scanner)
+	_, err := readMultilineString(t.Context(), strings.NewReader(input), "")
 	if err == nil {
-		t.Error("readMultilineInput() should return error for unclosed multiline input")
+		t.Fatal("readMultilineInput() should return an error for unclosed multiline input")
 	}
 
-	if !strings.Contains(err.Error(), "unclosed multiline input") {
-		t.Errorf("readMultilineInput() error should mention unclosed multiline input, got: %v", err)
-	}
+	assert.Contains(t, err.Error(), "unclosed multiline input", "error should mention unclosed multiline input")
+	// Error should also be io.EOF
+	assert.True(t, errors.Is(err, io.EOF), "error should be io.EOF")
 }

desktop/desktop_test.go

@@ -63,8 +63,9 @@ func TestChatHuggingFaceModel(t *testing.T) {
 		Body:       io.NopCloser(bytes.NewBufferString("data: {\"choices\":[{\"delta\":{\"content\":\"Hello there!\"}}]}\n")),
 	}, nil)
 
-	err := client.Chat("", modelName, prompt, "")
+	resp, err := client.Chat("", modelName, prompt, "")
 	assert.NoError(t, err)
+	assert.Equal(t, []string{"Hello there!"}, resp)
 }
 
 func TestInspectHuggingFaceModel(t *testing.T) {