1package shell
2
3import (
4 "context"
5 "strings"
6 "testing"
7 "time"
8)
9
10func TestBackgroundShellManager_Start(t *testing.T) {
11 t.Skip("Skipping this until I figure out why its flaky")
12 t.Parallel()
13
14 ctx := context.Background()
15 workingDir := t.TempDir()
16 manager := GetBackgroundShellManager()
17
18 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'hello world'", "")
19 if err != nil {
20 t.Fatalf("failed to start background shell: %v", err)
21 }
22
23 if bgShell.ID == "" {
24 t.Error("expected shell ID to be non-empty")
25 }
26
27 // Wait for the command to complete
28 bgShell.Wait()
29
30 stdout, stderr, done, err := bgShell.GetOutput()
31 if !done {
32 t.Error("expected shell to be done")
33 }
34
35 if err != nil {
36 t.Errorf("expected no error, got: %v", err)
37 }
38
39 if !strings.Contains(stdout, "hello world") {
40 t.Errorf("expected stdout to contain 'hello world', got: %s", stdout)
41 }
42
43 if stderr != "" {
44 t.Errorf("expected empty stderr, got: %s", stderr)
45 }
46}
47
48func TestBackgroundShellManager_Get(t *testing.T) {
49 t.Parallel()
50
51 ctx := context.Background()
52 workingDir := t.TempDir()
53 manager := GetBackgroundShellManager()
54
55 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'test'", "")
56 if err != nil {
57 t.Fatalf("failed to start background shell: %v", err)
58 }
59
60 // Retrieve the shell
61 retrieved, ok := manager.Get(bgShell.ID)
62 if !ok {
63 t.Error("expected to find the background shell")
64 }
65
66 if retrieved.ID != bgShell.ID {
67 t.Errorf("expected shell ID %s, got %s", bgShell.ID, retrieved.ID)
68 }
69
70 // Clean up
71 manager.Kill(bgShell.ID)
72}
73
74func TestBackgroundShellManager_Kill(t *testing.T) {
75 t.Parallel()
76
77 ctx := context.Background()
78 workingDir := t.TempDir()
79 manager := GetBackgroundShellManager()
80
81 // Start a long-running command
82 bgShell, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
83 if err != nil {
84 t.Fatalf("failed to start background shell: %v", err)
85 }
86
87 // Kill it
88 err = manager.Kill(bgShell.ID)
89 if err != nil {
90 t.Errorf("failed to kill background shell: %v", err)
91 }
92
93 // Verify it's no longer in the manager
94 _, ok := manager.Get(bgShell.ID)
95 if ok {
96 t.Error("expected shell to be removed after kill")
97 }
98
99 // Verify the shell is done
100 if !bgShell.IsDone() {
101 t.Error("expected shell to be done after kill")
102 }
103}
104
105func TestBackgroundShellManager_KillNonExistent(t *testing.T) {
106 t.Parallel()
107
108 manager := GetBackgroundShellManager()
109
110 err := manager.Kill("non-existent-id")
111 if err == nil {
112 t.Error("expected error when killing non-existent shell")
113 }
114}
115
116func TestBackgroundShell_IsDone(t *testing.T) {
117 t.Parallel()
118
119 ctx := context.Background()
120 workingDir := t.TempDir()
121 manager := GetBackgroundShellManager()
122
123 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'quick'", "")
124 if err != nil {
125 t.Fatalf("failed to start background shell: %v", err)
126 }
127
128 // Wait a bit for the command to complete
129 time.Sleep(100 * time.Millisecond)
130
131 if !bgShell.IsDone() {
132 t.Error("expected shell to be done")
133 }
134
135 // Clean up
136 manager.Kill(bgShell.ID)
137}
138
139func TestBackgroundShell_WithBlockFuncs(t *testing.T) {
140 t.Parallel()
141
142 ctx := context.Background()
143 workingDir := t.TempDir()
144 manager := GetBackgroundShellManager()
145
146 blockFuncs := []BlockFunc{
147 CommandsBlocker([]string{"curl", "wget"}),
148 }
149
150 bgShell, err := manager.Start(ctx, workingDir, blockFuncs, "curl example.com", "")
151 if err != nil {
152 t.Fatalf("failed to start background shell: %v", err)
153 }
154
155 // Wait for the command to complete
156 bgShell.Wait()
157
158 stdout, stderr, done, execErr := bgShell.GetOutput()
159 if !done {
160 t.Error("expected shell to be done")
161 }
162
163 // The command should have been blocked
164 output := stdout + stderr
165 if !strings.Contains(output, "not allowed") && execErr == nil {
166 t.Errorf("expected command to be blocked, got stdout: %s, stderr: %s, err: %v", stdout, stderr, execErr)
167 }
168
169 // Clean up
170 manager.Kill(bgShell.ID)
171}
172
173func TestBackgroundShellManager_List(t *testing.T) {
174 t.Parallel()
175
176 ctx := context.Background()
177 workingDir := t.TempDir()
178 manager := GetBackgroundShellManager()
179
180 // Start two shells
181 bgShell1, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
182 if err != nil {
183 t.Fatalf("failed to start first background shell: %v", err)
184 }
185
186 bgShell2, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
187 if err != nil {
188 t.Fatalf("failed to start second background shell: %v", err)
189 }
190
191 ids := manager.List()
192
193 // Check that both shells are in the list
194 found1 := false
195 found2 := false
196 for _, id := range ids {
197 if id == bgShell1.ID {
198 found1 = true
199 }
200 if id == bgShell2.ID {
201 found2 = true
202 }
203 }
204
205 if !found1 {
206 t.Errorf("expected to find shell %s in list", bgShell1.ID)
207 }
208 if !found2 {
209 t.Errorf("expected to find shell %s in list", bgShell2.ID)
210 }
211
212 // Clean up
213 manager.Kill(bgShell1.ID)
214 manager.Kill(bgShell2.ID)
215}
216
217func TestBackgroundShellManager_KillAll(t *testing.T) {
218 t.Parallel()
219
220 ctx := context.Background()
221 workingDir := t.TempDir()
222 manager := GetBackgroundShellManager()
223
224 // Start multiple long-running shells
225 shell1, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
226 if err != nil {
227 t.Fatalf("failed to start shell 1: %v", err)
228 }
229
230 shell2, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
231 if err != nil {
232 t.Fatalf("failed to start shell 2: %v", err)
233 }
234
235 shell3, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
236 if err != nil {
237 t.Fatalf("failed to start shell 3: %v", err)
238 }
239
240 // Verify shells are running
241 if shell1.IsDone() || shell2.IsDone() || shell3.IsDone() {
242 t.Error("shells should not be done yet")
243 }
244
245 // Kill all shells
246 manager.KillAll()
247
248 // Verify all shells are done
249 if !shell1.IsDone() {
250 t.Error("shell1 should be done after KillAll")
251 }
252 if !shell2.IsDone() {
253 t.Error("shell2 should be done after KillAll")
254 }
255 if !shell3.IsDone() {
256 t.Error("shell3 should be done after KillAll")
257 }
258
259 // Verify they're removed from the manager
260 if _, ok := manager.Get(shell1.ID); ok {
261 t.Error("shell1 should be removed from manager")
262 }
263 if _, ok := manager.Get(shell2.ID); ok {
264 t.Error("shell2 should be removed from manager")
265 }
266 if _, ok := manager.Get(shell3.ID); ok {
267 t.Error("shell3 should be removed from manager")
268 }
269
270 // Verify list is empty (or doesn't contain our shells)
271 ids := manager.List()
272 for _, id := range ids {
273 if id == shell1.ID || id == shell2.ID || id == shell3.ID {
274 t.Errorf("shell %s should not be in list after KillAll", id)
275 }
276 }
277}