background_test.go

  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}