background_test.go

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