1package modes
2
3import (
4 "context"
5 "runtime"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/patriceckhart/zot/packages/core"
11)
12
13func TestShellEscapeAddsOutputToAgentContext(t *testing.T) {
14 agent := core.NewAgent(nil, "", "", nil)
15 i := NewInteractive(InteractiveConfig{Agent: agent, CWD: t.TempDir()})
16 cmd := "printf zot-shell-context"
17 if runtime.GOOS == "windows" {
18 cmd = "echo zot-shell-context"
19 }
20
21 i.startShellEscape(context.Background(), cmd)
22
23 deadline := time.Now().Add(5 * time.Second)
24 for time.Now().Before(deadline) {
25 messages := agent.Messages()
26 i.mu.Lock()
27 running := i.shellRunning
28 i.mu.Unlock()
29 if len(messages) == 1 && !running {
30 text := userMessageText(messages[0])
31 if !strings.Contains(text, "zot-shell-context") {
32 t.Fatalf("shell context = %q, want command output", text)
33 }
34 if messages[0].Meta[shellEscapeMetaKey] != "true" {
35 t.Fatalf("shell context metadata = %v", messages[0].Meta)
36 }
37 history := i.inputHistory()
38 if len(history) != 1 || history[0] != "!"+cmd {
39 t.Fatalf("input history = %q, want shell command with bang", history)
40 }
41 return
42 }
43 time.Sleep(10 * time.Millisecond)
44 }
45 t.Fatal("shell output was not appended to agent context")
46}
47
48func TestShellEscapeCommand(t *testing.T) {
49 cases := []struct {
50 in string
51 wantCmd string
52 wantOK bool
53 }{
54 {"!ls -la", "ls -la", true},
55 {" !pwd", "pwd", true},
56 {"! go test ./... ", "go test ./...", true},
57 {"!", "", false},
58 {"! ", "", false},
59 {"ls -la", "", false},
60 {"/help", "", false},
61 {"hello !world", "", false},
62 }
63 for _, c := range cases {
64 cmd, ok := shellEscapeCommand(c.in)
65 if ok != c.wantOK || cmd != c.wantCmd {
66 t.Errorf("shellEscapeCommand(%q) = (%q,%v); want (%q,%v)",
67 c.in, cmd, ok, c.wantCmd, c.wantOK)
68 }
69 }
70}