1package shell
2
3import (
4 "context"
5 "runtime"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/stretchr/testify/require"
11)
12
13func TestBackgroundShellManager_Start(t *testing.T) {
14 t.Skip("Skipping this until I figure out why its flaky")
15 t.Parallel()
16
17 ctx := t.Context()
18 workingDir := t.TempDir()
19 manager := newBackgroundShellManager()
20
21 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'hello world'", "")
22 if err != nil {
23 t.Fatalf("failed to start background shell: %v", err)
24 }
25
26 if bgShell.ID == "" {
27 t.Error("expected shell ID to be non-empty")
28 }
29
30 // Wait for the command to complete
31 bgShell.Wait()
32
33 stdout, stderr, done, err := bgShell.GetOutput()
34 if !done {
35 t.Error("expected shell to be done")
36 }
37
38 if err != nil {
39 t.Errorf("expected no error, got: %v", err)
40 }
41
42 if !strings.Contains(stdout, "hello world") {
43 t.Errorf("expected stdout to contain 'hello world', got: %s", stdout)
44 }
45
46 if stderr != "" {
47 t.Errorf("expected empty stderr, got: %s", stderr)
48 }
49}
50
51func TestBackgroundShellManager_Get(t *testing.T) {
52 t.Parallel()
53
54 ctx := t.Context()
55 workingDir := t.TempDir()
56 manager := newBackgroundShellManager()
57
58 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'test'", "")
59 if err != nil {
60 t.Fatalf("failed to start background shell: %v", err)
61 }
62
63 // Retrieve the shell
64 retrieved, ok := manager.Get(bgShell.ID)
65 if !ok {
66 t.Error("expected to find the background shell")
67 }
68
69 if retrieved.ID != bgShell.ID {
70 t.Errorf("expected shell ID %s, got %s", bgShell.ID, retrieved.ID)
71 }
72
73 // Clean up
74 manager.Kill(bgShell.ID)
75}
76
77func TestBackgroundShellManager_Kill(t *testing.T) {
78 t.Parallel()
79
80 ctx := t.Context()
81 workingDir := t.TempDir()
82 manager := newBackgroundShellManager()
83
84 // Start a long-running command
85 bgShell, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
86 if err != nil {
87 t.Fatalf("failed to start background shell: %v", err)
88 }
89
90 // Kill it
91 err = manager.Kill(bgShell.ID)
92 if err != nil {
93 t.Errorf("failed to kill background shell: %v", err)
94 }
95
96 // Verify it's no longer in the manager
97 _, ok := manager.Get(bgShell.ID)
98 if ok {
99 t.Error("expected shell to be removed after kill")
100 }
101
102 // Verify the shell is done
103 if !bgShell.IsDone() {
104 t.Error("expected shell to be done after kill")
105 }
106}
107
108func TestBackgroundShellManager_KillNonExistent(t *testing.T) {
109 t.Parallel()
110
111 manager := newBackgroundShellManager()
112
113 err := manager.Kill("non-existent-id")
114 if err == nil {
115 t.Error("expected error when killing non-existent shell")
116 }
117}
118
119func TestBackgroundShell_IsDone(t *testing.T) {
120 t.Parallel()
121
122 ctx := t.Context()
123 workingDir := t.TempDir()
124 manager := newBackgroundShellManager()
125
126 bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'quick'", "")
127 if err != nil {
128 t.Fatalf("failed to start background shell: %v", err)
129 }
130
131 // Wait for the command to complete (Windows is slower to spin up).
132 require.Eventually(t, bgShell.IsDone, 5*time.Second, 50*time.Millisecond, "expected shell to be done")
133
134 // Clean up
135 manager.Kill(bgShell.ID)
136}
137
138func TestBackgroundShell_WithBlockFuncs(t *testing.T) {
139 t.Parallel()
140
141 ctx := t.Context()
142 workingDir := t.TempDir()
143 manager := newBackgroundShellManager()
144
145 blockFuncs := []BlockFunc{
146 CommandsBlocker([]string{"curl", "wget"}),
147 }
148
149 bgShell, err := manager.Start(ctx, workingDir, blockFuncs, "curl example.com", "")
150 if err != nil {
151 t.Fatalf("failed to start background shell: %v", err)
152 }
153
154 // Wait for the command to complete
155 bgShell.Wait()
156
157 stdout, stderr, done, execErr := bgShell.GetOutput()
158 if !done {
159 t.Error("expected shell to be done")
160 }
161
162 // The command should have been blocked
163 output := stdout + stderr
164 if !strings.Contains(output, "not allowed") && execErr == nil {
165 t.Errorf("expected command to be blocked, got stdout: %s, stderr: %s, err: %v", stdout, stderr, execErr)
166 }
167
168 // Clean up
169 manager.Kill(bgShell.ID)
170}
171
172func TestBackgroundShellManager_List(t *testing.T) {
173 if runtime.GOOS == "windows" {
174 t.Skip("skipping flacky test on windows")
175 }
176
177 t.Parallel()
178
179 ctx := t.Context()
180 workingDir := t.TempDir()
181 manager := newBackgroundShellManager()
182
183 // Start two shells
184 bgShell1, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
185 if err != nil {
186 t.Fatalf("failed to start first background shell: %v", err)
187 }
188
189 bgShell2, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
190 if err != nil {
191 t.Fatalf("failed to start second background shell: %v", err)
192 }
193
194 ids := manager.List()
195
196 // Check that both shells are in the list
197 found1 := false
198 found2 := false
199 for _, id := range ids {
200 if id == bgShell1.ID {
201 found1 = true
202 }
203 if id == bgShell2.ID {
204 found2 = true
205 }
206 }
207
208 if !found1 {
209 t.Errorf("expected to find shell %s in list", bgShell1.ID)
210 }
211 if !found2 {
212 t.Errorf("expected to find shell %s in list", bgShell2.ID)
213 }
214
215 // Clean up
216 manager.Kill(bgShell1.ID)
217 manager.Kill(bgShell2.ID)
218}
219
220func TestBackgroundShellManager_KillAll(t *testing.T) {
221 t.Parallel()
222
223 ctx := t.Context()
224 workingDir := t.TempDir()
225 manager := newBackgroundShellManager()
226
227 // Start multiple long-running shells
228 shell1, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
229 if err != nil {
230 t.Fatalf("failed to start shell 1: %v", err)
231 }
232
233 shell2, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
234 if err != nil {
235 t.Fatalf("failed to start shell 2: %v", err)
236 }
237
238 shell3, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
239 if err != nil {
240 t.Fatalf("failed to start shell 3: %v", err)
241 }
242
243 // Verify shells are running
244 if shell1.IsDone() || shell2.IsDone() || shell3.IsDone() {
245 t.Error("shells should not be done yet")
246 }
247
248 // Kill all shells
249 manager.KillAll(t.Context())
250
251 // Verify all shells are done
252 if !shell1.IsDone() {
253 t.Error("shell1 should be done after KillAll")
254 }
255 if !shell2.IsDone() {
256 t.Error("shell2 should be done after KillAll")
257 }
258 if !shell3.IsDone() {
259 t.Error("shell3 should be done after KillAll")
260 }
261
262 // Verify they're removed from the manager
263 if _, ok := manager.Get(shell1.ID); ok {
264 t.Error("shell1 should be removed from manager")
265 }
266 if _, ok := manager.Get(shell2.ID); ok {
267 t.Error("shell2 should be removed from manager")
268 }
269 if _, ok := manager.Get(shell3.ID); ok {
270 t.Error("shell3 should be removed from manager")
271 }
272
273 // Verify list is empty (or doesn't contain our shells)
274 ids := manager.List()
275 for _, id := range ids {
276 if id == shell1.ID || id == shell2.ID || id == shell3.ID {
277 t.Errorf("shell %s should not be in list after KillAll", id)
278 }
279 }
280}
281
282func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {
283 t.Parallel()
284
285 // XXX: can't use synctest here - causes --race to trip.
286
287 workingDir := t.TempDir()
288 manager := newBackgroundShellManager()
289
290 // Start a shell that traps signals and ignores cancellation.
291 _, err := manager.Start(t.Context(), workingDir, nil, "trap '' TERM INT; sleep 60", "")
292 require.NoError(t, err)
293
294 // Short timeout to test the timeout path.
295 ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
296 t.Cleanup(cancel)
297
298 start := time.Now()
299 manager.KillAll(ctx)
300
301 elapsed := time.Since(start)
302
303 // Must return promptly after timeout, not hang for 60 seconds.
304 require.Less(t, elapsed, 2*time.Second)
305}
306
307func TestBackgroundShell_WaitContext_Completed(t *testing.T) {
308 t.Parallel()
309
310 done := make(chan struct{})
311 close(done)
312
313 bgShell := &BackgroundShell{done: done}
314
315 ctx, cancel := context.WithTimeout(t.Context(), time.Second)
316 t.Cleanup(cancel)
317
318 require.True(t, bgShell.WaitContext(ctx))
319}
320
321func TestBackgroundShell_WaitContext_Canceled(t *testing.T) {
322 t.Parallel()
323
324 bgShell := &BackgroundShell{done: make(chan struct{})}
325
326 ctx, cancel := context.WithCancel(t.Context())
327 cancel()
328
329 require.False(t, bgShell.WaitContext(ctx))
330}