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