gosh/prompt_test.go

package main

import (
	"os"
	"strings"
	"testing"
	"time"
)

// unsetPS1 removes $PS1 for the test, so that the default prompt is used.
func unsetPS1(t *testing.T) {
	t.Helper()
	t.Setenv("PS1", "")
	os.Unsetenv("PS1")
}

var testInfo = promptInfo{
	user: "flow", host: "box.example.com", wd: "/home/flow/src/gosh", home: "/home/flow",
	jobs: 2, tty: "pts/3", now: time.Date(2026, 10, 5, 14, 7, 9, 0, time.UTC),
}

func TestExpandPrompt_Escapes(t *testing.T) {
	t.Setenv("GOSH_PS_VAR", "val")
	rootInfo := testInfo
	rootInfo.root, rootInfo.wd = true, "/"
	homeInfo := testInfo
	homeInfo.wd = "/home/flow"
	tests := []struct {
		ps   string
		info promptInfo
		want string
	}{
		{`\u@\h:\w\$ `, testInfo, "flow@box:~/src/gosh$ "},
		{`\u@\H \W \$ `, rootInfo, "flow@box.example.com / # "},
		{`\W`, homeInfo, "~"},
		{`\W`, testInfo, "gosh"},
		{`\t \T \@ \A \d`, testInfo, "14:07:09 02:07:09 02:07 PM 14:07 Mon Oct 05"},
		{`\j \l \s`, testInfo, "2 pts/3 gosh"},
		{`a\nb\\c`, testInfo, "a\nb\\c"},
		{`\[\e[1;32m\]x\[\033[0m\]`, testInfo, "\x1b[1;32mx\x1b[0m"},
		{`\101\400`, testInfo, "A 0"},
		{`\q\`, testInfo, `\q\`},
		{`$GOSH_PS_VAR ${GOSH_PS_VAR}-$((1+2)) $(echo sub) $?`, testInfo, "val val-3 sub 0"},
		{`'$GOSH_PS_VAR' "x"`, testInfo, `'val' "x"`},
		{`cost: 5$ \$`, testInfo, "cost: 5$ $"},
	}
	for _, tt := range tests {
		got, err := expandPrompt(tt.ps, tt.info)
		if err != nil {
			t.Errorf("expandPrompt(%q): %v", tt.ps, err)
		} else if got != tt.want {
			t.Errorf("expandPrompt(%q) = %q, want %q", tt.ps, got, tt.want)
		}
	}
}

func TestExpandPrompt_Errors(t *testing.T) {
	os.Unsetenv("GOSH_PS_UNSET")
	if _, err := expandPrompt(`$GOSH_PS_UNSET`, testInfo); err == nil {
		t.Error("expected an error for an unset variable")
	}
	if got, err := expandPrompt(`${GOSH_PS_UNSET:-none}`, testInfo); err != nil || got != "none" {
		t.Errorf("got %q, %v", got, err)
	}
}

func TestExpandPrompt_KeepsStatus(t *testing.T) {
	lastStatus, substStatus = 3, 0
	defer func() { lastStatus = 0 }()
	got, err := expandPrompt(`$(exit 7)$?`, testInfo)
	if err != nil || got != "3" {
		t.Errorf("got %q, %v", got, err)
	}
	if substStatus != 0 {
		t.Errorf("substStatus = %d, want 0", substStatus)
	}
}

func TestExpandPrompt_Sanitizes(t *testing.T) {
	t.Setenv("GOSH_PS_EVIL", "a\x1b[2Jb")
	info := testInfo
	info.wd = "/tmp/evil\x1b]0;x\x07"
	got, err := expandPrompt(`\w $GOSH_PS_EVIL $(printf 'c\033[2Jd')`, info)
	if err != nil {
		t.Fatal(err)
	}
	if strings.ContainsAny(got, "\x1b\x07") {
		t.Errorf("prompt contains control characters: %q", got)
	}
}

func TestSplitPrompt(t *testing.T) {
	tests := []struct{ prompt, prefix, line string }{
		{"a $ ", "", "a $ "},
		{"top\n\x1b[1mbold\x1b[0m $ ", "top\n", "\x1b[1mbold\x1b[0m $ "},
		{"l1\nl2\n$ ", "l1\nl2\n", "$ "},
		{"\x1b]0;title\x07$ ", "\x1b]0;title\x07", "$ "},
		{"\x1b]0;t\x1b\\\x1b[2K$ \a", "\x1b]0;t\x1b\\\x1b[2K\a", "$ "},
		{"x\x1b", "\x1b", "x"},
		{"x\x1b[", "\x1b[", "x"},
	}
	for _, tt := range tests {
		prefix, line := splitPrompt(tt.prompt)
		if prefix != tt.prefix || line != tt.line {
			t.Errorf("splitPrompt(%q) = %q, %q, want %q, %q", tt.prompt, prefix, line, tt.prefix, tt.line)
		}
	}
}

func TestGetPrompt_PS1(t *testing.T) {
	t.Setenv("PS1", `[\s] \$ `)
	want := "[gosh] $ "
	if os.Geteuid() == 0 {
		want = "[gosh] # "
	}
	if got := getPrompt(); got != want {
		t.Errorf("getPrompt() = %q, want %q", got, want)
	}
	// An invalid PS1 falls back to the default prompt
	t.Setenv("PS1", `$GOSH_PS_UNSET`)
	os.Unsetenv("GOSH_PS_UNSET")
	if got := getPrompt(); !strings.Contains(got, "@") {
		t.Errorf("expected default prompt, got %q", got)
	}
}