loveuer
72 lines
1.8 KiB
Go
72 lines
1.8 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
)
|
|
|
|
func TestShellRunStreamsOutputWhileRunning(t *testing.T) {
|
|
tool := NewShellRunTool(t.TempDir())
|
|
var mu sync.Mutex
|
|
var chunks []string
|
|
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"printf 'a\\n'; sleep 0.05; printf 'b\\n'"}`), func(chunk string) error {
|
|
mu.Lock()
|
|
chunks = append(chunks, chunk)
|
|
mu.Unlock()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.Contains(output, "a") || !strings.Contains(output, "b") {
|
|
t.Fatalf("output = %q", output)
|
|
}
|
|
mu.Lock()
|
|
joined := strings.Join(chunks, "")
|
|
mu.Unlock()
|
|
if len(chunks) == 0 {
|
|
t.Fatal("expected at least one streamed chunk")
|
|
}
|
|
if !strings.Contains(joined, "a") || !strings.Contains(joined, "b") {
|
|
t.Fatalf("streamed chunks = %q", joined)
|
|
}
|
|
}
|
|
|
|
func TestShellRunExecuteCollectsFullOutput(t *testing.T) {
|
|
tool := NewShellRunTool(t.TempDir())
|
|
output, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"printf 'hello world'"}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.TrimSpace(output) != "hello world" {
|
|
t.Fatalf("output = %q", output)
|
|
}
|
|
}
|
|
|
|
func TestShellRunStreamsStderrAndStdout(t *testing.T) {
|
|
tool := NewShellRunTool(t.TempDir())
|
|
var mu sync.Mutex
|
|
var chunks []string
|
|
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"echo out; echo err >&2"}`), func(chunk string) error {
|
|
mu.Lock()
|
|
chunks = append(chunks, chunk)
|
|
mu.Unlock()
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !strings.Contains(output, "out") || !strings.Contains(output, "err") {
|
|
t.Fatalf("output = %q", output)
|
|
}
|
|
mu.Lock()
|
|
joined := strings.Join(chunks, "")
|
|
mu.Unlock()
|
|
if !strings.Contains(joined, "out") || !strings.Contains(joined, "err") {
|
|
t.Fatalf("streamed chunks = %q", joined)
|
|
}
|
|
}
|