Vectrum
Sign In
267 lines · 8.7 KB · Go
package tools

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

// waitFor polls until cond holds, so a test never sleeps longer than it must
// and never races a process that has not been scheduled yet.
func waitFor(t *testing.T, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("timed out waiting for %s", what)
}

// A dev server is the case this exists for: the call must come back with a
// handle immediately rather than blocking until the server stops.
func TestBackgroundBashReturnsAHandleWithoutWaiting(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()

	start := time.Now()
	out, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command:    "echo started; sleep 30",
		Background: true,
	})
	if err != nil {
		t.Fatalf("background run: %v", err)
	}
	if elapsed := time.Since(start); elapsed > 5*time.Second {
		t.Errorf("the call blocked for %s; a background command must return at once", elapsed)
	}
	if !strings.Contains(out, "bg_1") {
		t.Errorf("result %q does not name the handle the model has to poll with", out)
	}
	if !strings.Contains(out, "process_output") {
		t.Errorf("result %q does not say how to read the output", out)
	}
}

// Output accumulates while the process runs, and each read returns only what
// is new: polling a chatty server must not re-send its whole log every turn.
func TestProcessOutputIsIncremental(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	dir := t.TempDir()

	if _, err := run(t, NewBash(dir, WithProcesses(procs)), BashArgs{
		Command:    "echo first; sleep 0.4; echo second; sleep 30",
		Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	outTool := NewProcessOutput(procs)

	waitFor(t, "the first line", func() bool {
		got, _ := run(t, outTool, ProcessOutputArgs{ID: "bg_1"})
		return strings.Contains(body(got), "first")
	})
	// "first" has been handed over, so it must not come back a second time.
	waitFor(t, "the second line", func() bool {
		got, _ := run(t, outTool, ProcessOutputArgs{ID: "bg_1"})
		if strings.Contains(body(got), "first") {
			t.Errorf("output was re-sent after being read:\n%s", got)
		}
		return strings.Contains(body(got), "second")
	})
}

// body is everything after the status line. The status line quotes the command
// itself, so a test looking for a word the command also contains has to skip
// it or it matches the echo rather than the output.
func body(report string) string {
	_, rest, _ := strings.Cut(report, "\n")
	return rest
}

// Status is what the model acts on: still running, or finished with a code.
func TestProcessOutputReportsExit(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()

	if _, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command: "echo done; exit 3", Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	outTool := NewProcessOutput(procs)
	waitFor(t, "the process to exit", func() bool {
		got, _ := run(t, outTool, ProcessOutputArgs{ID: "bg_1"})
		return strings.Contains(got, "exit code: 3")
	})
	got, err := run(t, outTool, ProcessOutputArgs{ID: "bg_1"})
	if err != nil {
		t.Fatalf("reading a finished process is not an error: %v", err)
	}
	if !strings.Contains(got, "exit code: 3") {
		t.Errorf("a finished process must keep reporting how it ended, got %q", got)
	}
}

// Killing must take the children with it. A dev server started from a shell
// is a child of that shell, and killing only the shell leaves the port held.
func TestProcessKillTakesTheWholeGroup(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	dir := t.TempDir()
	marker := filepath.Join(dir, "child-survived")

	if _, err := run(t, NewBash(dir, WithProcesses(procs)), BashArgs{
		// The inner shell outlives its parent unless the group is killed.
		Command:    "(sleep 1; touch " + marker + ") & sleep 30",
		Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	waitFor(t, "the process to start", func() bool {
		got, _ := run(t, NewProcessOutput(procs), ProcessOutputArgs{ID: "bg_1"})
		return strings.Contains(got, "running")
	})

	if _, err := run(t, NewProcessKill(procs), ProcessKillArgs{ID: "bg_1"}); err != nil {
		t.Fatalf("kill: %v", err)
	}
	time.Sleep(1500 * time.Millisecond)
	if _, err := os.Stat(marker); err == nil {
		t.Error("a child of the killed process was still alive and finished its work")
	}
}

// An id the model made up costs a turn if the error says nothing useful.
func TestUnknownProcessIDNamesTheRealOnes(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	if _, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command: "sleep 30", Background: true,
	}); err != nil {
		t.Fatal(err)
	}

	_, err := run(t, NewProcessOutput(procs), ProcessOutputArgs{ID: "bg_99"})
	if err == nil {
		t.Fatal("an unknown id must be an error")
	}
	if !strings.Contains(err.Error(), "bg_1") {
		t.Errorf("error %q should list the ids that do exist", err)
	}
}

// With no id at all the tool lists what is running, so the model can recover
// from having lost track without guessing.
func TestProcessOutputWithoutIDLists(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	if _, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command: "sleep 30", Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	got, err := run(t, NewProcessOutput(procs), ProcessOutputArgs{})
	if err != nil {
		t.Fatalf("listing: %v", err)
	}
	if !strings.Contains(got, "bg_1") || !strings.Contains(got, "sleep 30") {
		t.Errorf("listing %q should name each process and the command behind it", got)
	}
}

// Nothing may outlive the session: an orphaned dev server holds its port and
// the user never sees the process that is doing it.
func TestCloseKillsEverything(t *testing.T) {
	procs := NewProcesses()
	dir := t.TempDir()
	marker := filepath.Join(dir, "survived")
	if _, err := run(t, NewBash(dir, WithProcesses(procs)), BashArgs{
		Command: "sleep 0.8; touch " + marker, Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	waitFor(t, "the process to start", func() bool {
		got, _ := run(t, NewProcessOutput(procs), ProcessOutputArgs{ID: "bg_1"})
		return strings.Contains(got, "running")
	})

	procs.Close()
	time.Sleep(1200 * time.Millisecond)
	if _, err := os.Stat(marker); err == nil {
		t.Error("a background process survived Close and finished its work")
	}
}

// A server that logs forever must not grow the buffer without bound.
func TestBackgroundOutputIsCapped(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	if _, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command: "seq 1 200000", Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	outTool := NewProcessOutput(procs)
	waitFor(t, "the process to finish", func() bool {
		got, _ := run(t, outTool, ProcessOutputArgs{ID: "bg_1"})
		return strings.Contains(got, "exit code: 0")
	})
	if n := procs.buffered("bg_1"); n > maxBackgroundOutput {
		t.Errorf("buffered %d bytes, want at most %d", n, maxBackgroundOutput)
	}
}

// Without the option the tool is exactly what it was: background is refused
// with an error that says so rather than silently running in the foreground.
func TestBackgroundRefusedWhenUnsupported(t *testing.T) {
	_, err := run(t, NewBash(t.TempDir()), BashArgs{Command: "echo hi", Background: true})
	if err == nil || !strings.Contains(err.Error(), "background") {
		t.Errorf("err = %v, want a refusal naming the unsupported feature", err)
	}
}

func TestForegroundStillWorksWithProcesses(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	out, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{Command: "echo hello"})
	if err != nil {
		t.Fatalf("foreground: %v", err)
	}
	if !strings.Contains(out, "hello") || !strings.Contains(out, "exit code: 0") {
		t.Errorf("foreground output changed shape: %q", out)
	}
}

func TestKillIsIdempotent(t *testing.T) {
	procs := NewProcesses()
	defer procs.Close()
	if _, err := run(t, NewBash(t.TempDir(), WithProcesses(procs)), BashArgs{
		Command: "sleep 30", Background: true,
	}); err != nil {
		t.Fatal(err)
	}
	kill := NewProcessKill(procs)
	if _, err := run(t, kill, ProcessKillArgs{ID: "bg_1"}); err != nil {
		t.Fatalf("first kill: %v", err)
	}
	// Killing something already dead is a no-op, not an error: the model
	// cleaning up at the end of a task should not have to check first.
	if _, err := run(t, kill, ProcessKillArgs{ID: "bg_1"}); err != nil {
		t.Errorf("second kill: %v", err)
	}
}

func TestProcessesToolsAreNilWithoutAManager(t *testing.T) {
	if NewProcessOutput(nil) != nil || NewProcessKill(nil) != nil {
		t.Error("with no manager the tools must not be registered at all")
	}
}

var _ = context.Background