Vectrum
Sign In
272 lines · 8.9 KB · Go
package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func run(t *testing.T, tool Tool, args any) (string, error) {
	t.Helper()
	raw, err := json.Marshal(args)
	if err != nil {
		t.Fatal(err)
	}
	return tool.Execute(context.Background(), raw)
}

func TestBashSuccess(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "hello.txt", "hi\n")
	out, err := run(t, NewBash(dir), BashArgs{Command: "cat hello.txt"})
	if err != nil {
		t.Fatalf("bash: %v", err)
	}
	if !strings.Contains(out, "exit code: 0") || !strings.Contains(out, "--- stdout ---\nhi") {
		t.Errorf("got:\n%s", out)
	}
	if strings.Contains(out, "stderr") {
		t.Errorf("empty streams should be omitted:\n%s", out)
	}
}

func TestBashNonZeroExitIsAnErrorCarryingOutput(t *testing.T) {
	_, err := run(t, NewBash(t.TempDir()), BashArgs{Command: "echo out; echo problem >&2; exit 3"})
	if err == nil {
		t.Fatal("a failing command must come back as an error the model reacts to")
	}
	for _, want := range []string{"exit code: 3", "--- stdout ---\nout", "--- stderr ---\nproblem"} {
		if !strings.Contains(err.Error(), want) {
			t.Errorf("missing %q in:\n%s", want, err)
		}
	}
}

func TestBashCapsOutputKeepingBothEnds(t *testing.T) {
	out, err := run(t, NewBash(t.TempDir()), BashArgs{Command: "seq 1 5000"})
	if err != nil {
		t.Fatalf("bash: %v", err)
	}
	if !strings.Contains(out, "\n1\n") {
		t.Error("head was dropped")
	}
	if !strings.Contains(out, "5000") {
		t.Error("tail was dropped — that is where test failures live")
	}
	if !strings.Contains(out, "lines elided") {
		t.Error("truncation should be announced")
	}
	if n := strings.Count(out, "\n"); n > bashHeadLines+bashTailLines+10 {
		t.Errorf("output not capped: %d lines", n)
	}
}

func TestBashTimeoutKillsTheProcessGroup(t *testing.T) {
	if testing.Short() {
		t.Skip("waits on a real child process")
	}
	dir := t.TempDir()
	marker := filepath.Join(dir, "child-survived")
	// The child outlives its parent shell unless the whole group is killed.
	cmd := fmt.Sprintf(`(sleep 2; touch %q) & sleep 5`, marker)

	start := time.Now()
	_, err := run(t, NewBash(dir), BashArgs{Command: cmd, TimeoutMS: 300})
	if err == nil || !strings.Contains(err.Error(), "timed out") {
		t.Fatalf("error = %v, want a timeout", err)
	}
	if elapsed := time.Since(start); elapsed > 3*time.Second {
		t.Errorf("took %s: the timeout did not take effect promptly", elapsed)
	}

	time.Sleep(3 * time.Second)
	if _, err := os.Stat(marker); err == nil {
		t.Error("the orphaned child kept running: the process group was not killed")
	}
}

func TestBashRunsInTheWorkingRoot(t *testing.T) {
	dir := t.TempDir()
	out, err := run(t, NewBash(dir), BashArgs{Command: "pwd"})
	if err != nil {
		t.Fatalf("bash: %v", err)
	}
	resolved, _ := filepath.EvalSymlinks(dir)
	if !strings.Contains(out, dir) && !strings.Contains(out, resolved) {
		t.Errorf("pwd = %s, want the working root %s", out, dir)
	}
}

func TestBashRespectsCancellation(t *testing.T) {
	raw, _ := json.Marshal(BashArgs{Command: "sleep 10"})
	ctx, cancel := context.WithCancel(context.Background())
	go func() { time.Sleep(100 * time.Millisecond); cancel() }()

	start := time.Now()
	if _, err := NewBash(t.TempDir()).Execute(ctx, raw); err == nil {
		t.Fatal("want an error")
	}
	if elapsed := time.Since(start); elapsed > 3*time.Second {
		t.Errorf("took %s: cancellation did not propagate", elapsed)
	}
}

func TestBashEmptyCommand(t *testing.T) {
	if _, err := run(t, NewBash(t.TempDir()), BashArgs{Command: "   "}); err == nil ||
		!strings.Contains(err.Error(), "command is required") {
		t.Errorf("error = %v", err)
	}
}

func TestWriteFileCreatesIncludingParents(t *testing.T) {
	dir := t.TempDir()
	out, err := run(t, NewWrite(dir, newTracker()), WriteArgs{
		Path: "internal/new/thing.go", Content: "package thing\n",
	})
	if err != nil {
		t.Fatalf("write: %v", err)
	}
	got, err := os.ReadFile(filepath.Join(dir, "internal/new/thing.go"))
	if err != nil {
		t.Fatalf("file not created: %v", err)
	}
	if string(got) != "package thing\n" {
		t.Errorf("content = %q", got)
	}
	if !strings.Contains(out, "Created") || !strings.Contains(out, "14 bytes") {
		t.Errorf("summary = %q", out)
	}
}

func TestWriteFileRefusesToClobberAnUnreadFile(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "existing.go", "package existing\n")

	_, err := run(t, NewWrite(dir, newTracker()), WriteArgs{Path: "existing.go", Content: "package gone\n"})
	if err == nil || !strings.Contains(err.Error(), "has not been read this session") {
		t.Fatalf("error = %v, want a read-before-write refusal", err)
	}
	if !strings.Contains(err.Error(), "replaces the whole file") {
		t.Errorf("error should say why that matters here: %v", err)
	}
	got, _ := os.ReadFile(filepath.Join(dir, "existing.go"))
	if string(got) != "package existing\n" {
		t.Error("file must be untouched")
	}
}

func TestWriteFileReplacesAReadFile(t *testing.T) {
	dir := t.TempDir()
	path := writeFile(t, dir, "existing.go", "package existing\n")
	track := newTracker()
	info, _ := os.Stat(path)
	track.MarkRead(path, info.ModTime())

	out, err := run(t, NewWrite(dir, track), WriteArgs{Path: "existing.go", Content: "package replaced\n"})
	if err != nil {
		t.Fatalf("write: %v", err)
	}
	if !strings.Contains(out, "Replaced") {
		t.Errorf("summary = %q, want it to distinguish a rewrite from a create", out)
	}
	got, _ := os.ReadFile(path)
	if string(got) != "package replaced\n" {
		t.Errorf("content = %q", got)
	}
}

func TestWriteFilePreservesCRLFOnRewrite(t *testing.T) {
	dir := t.TempDir()
	path := writeFile(t, dir, "win.txt", "one\r\ntwo\r\n")
	track := newTracker()
	info, _ := os.Stat(path)
	track.MarkRead(path, info.ModTime())

	if _, err := run(t, NewWrite(dir, track), WriteArgs{Path: "win.txt", Content: "three\nfour\n"}); err != nil {
		t.Fatalf("write: %v", err)
	}
	got, _ := os.ReadFile(path)
	if string(got) != "three\r\nfour\r\n" {
		t.Errorf("got %q, want CRLF preserved", got)
	}
}

func TestRunCommandReportsExitCodeAndBothStreams(t *testing.T) {
	report, failed := RunCommand(context.Background(), t.TempDir(),
		[]string{"bash", "-c", "echo out; echo problem >&2; exit 3"}, time.Minute)
	if !failed {
		t.Error("a non-zero exit is a failure the caller has to be able to see")
	}
	for _, want := range []string{"exit code: 3", "--- stdout ---\nout", "--- stderr ---\nproblem"} {
		if !strings.Contains(report, want) {
			t.Errorf("missing %q in:\n%s", want, report)
		}
	}
}

func TestRunCommandSucceedsInTheGivenDirectory(t *testing.T) {
	dir := t.TempDir()
	writeFile(t, dir, "hello.txt", "hi\n")
	report, failed := RunCommand(context.Background(), dir, []string{"bash", "-c", "cat hello.txt"}, time.Minute)
	if failed {
		t.Errorf("want success, got:\n%s", report)
	}
	if !strings.Contains(report, "--- stdout ---\nhi") {
		t.Errorf("got:\n%s", report)
	}
}

func TestRunCommandReportsAMissingBinaryRatherThanAnExitCode(t *testing.T) {
	report, failed := RunCommand(context.Background(), t.TempDir(),
		[]string{"definitely-not-a-real-binary-xyz"}, time.Minute)
	if !failed {
		t.Fatal("a binary that is not there is a failure")
	}
	if !strings.Contains(report, "failed to run") {
		t.Errorf("want a run failure, got:\n%s", report)
	}
}

// A command that greps the process table must not find itself. With
// `bash -c "<command>"` the whole script sits in the shell's own argv, so
// `pkill -f "vite preview"` killed the shell that was running it — twice in
// one session. Running from a file keeps the text out of every argv.
func TestBashCommandTextIsNotInTheShellsArgv(t *testing.T) {
	out, err := run(t, NewBash(t.TempDir()), BashArgs{
		Command: `pgrep -af "needle-in-argv-xyzzy" | grep -v pgrep | grep -v grep; echo "count=$(pgrep -fc "needle-in-argv-xyzzy" || true)"`,
	})
	if err != nil {
		t.Fatalf("run: %v", err)
	}
	// pgrep itself carries the pattern, and pgrep -c excludes itself, so a
	// clean run counts zero. One is the shell holding our command text.
	if !strings.Contains(out, "count=0") {
		t.Errorf("the command text is visible in some process's argv:\n%s", out)
	}
}

// A file that came into being during this session — a scaffold's output, a
// file the model made with bash — is not the user's work, and may be replaced
// without a read. A file that predates the session still needs one.
func TestWriteFileMayReplaceAFileCreatedThisSession(t *testing.T) {
	dir := t.TempDir()
	old := writeFile(t, dir, "old.go", "package old\n")
	past := time.Now().Add(-time.Hour)
	os.Chtimes(old, past, past)

	tool := NewWrite(dir, newTracker())
	writeFile(t, dir, "fresh.go", "package scaffold\n") // appears after the tool exists

	if _, err := run(t, tool, WriteArgs{Path: "fresh.go", Content: "package mine\n"}); err != nil {
		t.Errorf("a file created this session should be replaceable unread: %v", err)
	}
	if _, err := run(t, tool, WriteArgs{Path: "old.go", Content: "package gone\n"}); err == nil ||
		!strings.Contains(err.Error(), "has not been read") {
		t.Errorf("a file older than the session still needs a read, got %v", err)
	}
}