on 8ebab21 help sign in
crush/internal/shell/background_test.go
334 lines 7.7 kB history raw
1package shell
3import (
4 "context"
5 "runtime"
6 "strings"
7 "testing"
8 "time"
10 "github.com/stretchr/testify/require"
13func TestBackgroundShellManager_Start(t *testing.T) {
14 t.Skip("Skipping this until I figure out why its flaky")
15 t.Parallel()
17 ctx := t.Context()
18 workingDir := t.TempDir()
19 manager := newBackgroundShellManager()
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 }
26 if bgShell.ID == "" {
27 t.Error("expected shell ID to be non-empty")
28 }
30 // Wait for the command to complete
31 bgShell.Wait()
33 stdout, stderr, done, err := bgShell.GetOutput()
34 if !done {
35 t.Error("expected shell to be done")
36 }
38 if err != nil {
39 t.Errorf("expected no error, got: %v", err)
40 }
42 if !strings.Contains(stdout, "hello world") {
43 t.Errorf("expected stdout to contain 'hello world', got: %s", stdout)
44 }
46 if stderr != "" {
47 t.Errorf("expected empty stderr, got: %s", stderr)
48 }
51func TestBackgroundShellManager_Get(t *testing.T) {
52 t.Parallel()
54 ctx := t.Context()
55 workingDir := t.TempDir()
56 manager := newBackgroundShellManager()
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 }
63 // Retrieve the shell
64 retrieved, ok := manager.Get(bgShell.ID)
65 if !ok {
66 t.Error("expected to find the background shell")
67 }
69 if retrieved.ID != bgShell.ID {
70 t.Errorf("expected shell ID %s, got %s", bgShell.ID, retrieved.ID)
71 }
73 // Clean up
74 manager.Kill(bgShell.ID)
77func TestBackgroundShellManager_Kill(t *testing.T) {
78 t.Parallel()
80 ctx := t.Context()
81 workingDir := t.TempDir()
82 manager := newBackgroundShellManager()
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 }
90 // Kill it
91 err = manager.Kill(bgShell.ID)
92 if err != nil {
93 t.Errorf("failed to kill background shell: %v", err)
94 }
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 }
102 // Verify the shell is done
103 if !bgShell.IsDone() {
104 t.Error("expected shell to be done after kill")
105 }
108func TestBackgroundShellManager_KillNonExistent(t *testing.T) {
109 t.Parallel()
111 manager := newBackgroundShellManager()
113 err := manager.Kill("non-existent-id")
114 if err == nil {
115 t.Error("expected error when killing non-existent shell")
116 }
119func TestBackgroundShell_IsDone(t *testing.T) {
120 t.Parallel()
122 ctx := t.Context()
123 workingDir := t.TempDir()
124 manager := newBackgroundShellManager()
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 }
131 // Wait a bit for the command to complete
132 time.Sleep(100 * time.Millisecond)
134 if !bgShell.IsDone() {
135 t.Error("expected shell to be done")
136 }
138 // Clean up
139 manager.Kill(bgShell.ID)
142func TestBackgroundShell_WithBlockFuncs(t *testing.T) {
143 t.Parallel()
145 ctx := t.Context()
146 workingDir := t.TempDir()
147 manager := newBackgroundShellManager()
149 blockFuncs := []BlockFunc{
150 CommandsBlocker([]string{"curl", "wget"}),
151 }
153 bgShell, err := manager.Start(ctx, workingDir, blockFuncs, "curl example.com", "")
154 if err != nil {
155 t.Fatalf("failed to start background shell: %v", err)
156 }
158 // Wait for the command to complete
159 bgShell.Wait()
161 stdout, stderr, done, execErr := bgShell.GetOutput()
162 if !done {
163 t.Error("expected shell to be done")
164 }
166 // The command should have been blocked
167 output := stdout + stderr
168 if !strings.Contains(output, "not allowed") && execErr == nil {
169 t.Errorf("expected command to be blocked, got stdout: %s, stderr: %s, err: %v", stdout, stderr, execErr)
170 }
172 // Clean up
173 manager.Kill(bgShell.ID)
176func TestBackgroundShellManager_List(t *testing.T) {
177 if runtime.GOOS == "windows" {
178 t.Skip("skipping flacky test on windows")
179 }
181 t.Parallel()
183 ctx := t.Context()
184 workingDir := t.TempDir()
185 manager := newBackgroundShellManager()
187 // Start two shells
188 bgShell1, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
189 if err != nil {
190 t.Fatalf("failed to start first background shell: %v", err)
191 }
193 bgShell2, err := manager.Start(ctx, workingDir, nil, "sleep 1", "")
194 if err != nil {
195 t.Fatalf("failed to start second background shell: %v", err)
196 }
198 ids := manager.List()
200 // Check that both shells are in the list
201 found1 := false
202 found2 := false
203 for _, id := range ids {
204 if id == bgShell1.ID {
205 found1 = true
206 }
207 if id == bgShell2.ID {
208 found2 = true
209 }
210 }
212 if !found1 {
213 t.Errorf("expected to find shell %s in list", bgShell1.ID)
214 }
215 if !found2 {
216 t.Errorf("expected to find shell %s in list", bgShell2.ID)
217 }
219 // Clean up
220 manager.Kill(bgShell1.ID)
221 manager.Kill(bgShell2.ID)
224func TestBackgroundShellManager_KillAll(t *testing.T) {
225 t.Parallel()
227 ctx := t.Context()
228 workingDir := t.TempDir()
229 manager := newBackgroundShellManager()
231 // Start multiple long-running shells
232 shell1, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
233 if err != nil {
234 t.Fatalf("failed to start shell 1: %v", err)
235 }
237 shell2, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
238 if err != nil {
239 t.Fatalf("failed to start shell 2: %v", err)
240 }
242 shell3, err := manager.Start(ctx, workingDir, nil, "sleep 10", "")
243 if err != nil {
244 t.Fatalf("failed to start shell 3: %v", err)
245 }
247 // Verify shells are running
248 if shell1.IsDone() || shell2.IsDone() || shell3.IsDone() {
249 t.Error("shells should not be done yet")
250 }
252 // Kill all shells
253 manager.KillAll(t.Context())
255 // Verify all shells are done
256 if !shell1.IsDone() {
257 t.Error("shell1 should be done after KillAll")
258 }
259 if !shell2.IsDone() {
260 t.Error("shell2 should be done after KillAll")
261 }
262 if !shell3.IsDone() {
263 t.Error("shell3 should be done after KillAll")
264 }
266 // Verify they're removed from the manager
267 if _, ok := manager.Get(shell1.ID); ok {
268 t.Error("shell1 should be removed from manager")
269 }
270 if _, ok := manager.Get(shell2.ID); ok {
271 t.Error("shell2 should be removed from manager")
272 }
273 if _, ok := manager.Get(shell3.ID); ok {
274 t.Error("shell3 should be removed from manager")
275 }
277 // Verify list is empty (or doesn't contain our shells)
278 ids := manager.List()
279 for _, id := range ids {
280 if id == shell1.ID || id == shell2.ID || id == shell3.ID {
281 t.Errorf("shell %s should not be in list after KillAll", id)
282 }
283 }
286func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) {
287 t.Parallel()
289 // XXX: can't use synctest here - causes --race to trip.
291 workingDir := t.TempDir()
292 manager := newBackgroundShellManager()
294 // Start a shell that traps signals and ignores cancellation.
295 _, err := manager.Start(t.Context(), workingDir, nil, "trap '' TERM INT; sleep 60", "")
296 require.NoError(t, err)
298 // Short timeout to test the timeout path.
299 ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
300 t.Cleanup(cancel)
302 start := time.Now()
303 manager.KillAll(ctx)
305 elapsed := time.Since(start)
307 // Must return promptly after timeout, not hang for 60 seconds.
308 require.Less(t, elapsed, 2*time.Second)
311func TestBackgroundShell_WaitContext_Completed(t *testing.T) {
312 t.Parallel()
314 done := make(chan struct{})
315 close(done)
317 bgShell := &BackgroundShell{done: done}
319 ctx, cancel := context.WithTimeout(t.Context(), time.Second)
320 t.Cleanup(cancel)
322 require.True(t, bgShell.WaitContext(ctx))
325func TestBackgroundShell_WaitContext_Canceled(t *testing.T) {
326 t.Parallel()
328 bgShell := &BackgroundShell{done: make(chan struct{})}
330 ctx, cancel := context.WithCancel(t.Context())
331 cancel()
333 require.False(t, bgShell.WaitContext(ctx))