1package swarm
2
3import (
4 "context"
5 "errors"
6 "strings"
7 "sync"
8 "testing"
9 "time"
10)
11
12// newTestSwarm builds a Swarm rooted in t.TempDir with the in-memory
13// and a Runner factory the test controls. Returns the swarm plus a
14// slice of runners keyed by spawn order so tests can assert they
15// were actually invoked.
16func newTestSwarm(t *testing.T, mk func(a *Agent) Runner) *Swarm {
17 t.Helper()
18 root := t.TempDir()
19 return New(Config{
20 Root: root,
21 RepoRoot: root,
22 NewRunner: mk,
23 })
24}
25
26func TestSpawnRunsAndCompletes(t *testing.T) {
27 ran := make(chan string, 1)
28 f := newTestSwarm(t, func(a *Agent) Runner {
29 return RunnerFunc(func(ctx context.Context, sink Sink) error {
30 ran <- a.Task
31 sink.Activity("hello")
32 sink.Transcript("line one")
33 sink.Transcript("line two")
34 return nil
35 })
36 })
37 a, err := f.Spawn(context.Background(), "do a thing")
38 if err != nil {
39 t.Fatal(err)
40 }
41 select {
42 case got := <-ran:
43 if got != "do a thing" {
44 t.Fatalf("runner got task %q; want %q", got, "do a thing")
45 }
46 case <-time.After(time.Second):
47 t.Fatal("runner did not start")
48 }
49 a.Wait()
50 if a.Status() != StatusDone {
51 t.Fatalf("status %s; want done", a.Status())
52 }
53 if got := a.Transcript(); len(got) != 2 || got[0] != "line one" || got[1] != "line two" {
54 t.Fatalf("transcript = %q", got)
55 }
56 if !strings.Contains(a.ID, "do-a-thing") {
57 t.Fatalf("id %q missing slug", a.ID)
58 }
59 // Every agent shares the host's RepoRoot.
60 if a.Dir != f.cfg.RepoRoot {
61 t.Fatalf("dir = %q; want repo root %q", a.Dir, f.cfg.RepoRoot)
62 }
63}
64
65// TestSpawnAgentSharesRepoRoot verifies the only-mode-we-support:
66// every spawned agent points its cwd at the parent zot's RepoRoot.
67func TestSpawnAgentSharesRepoRoot(t *testing.T) {
68 f := newTestSwarm(t, func(a *Agent) Runner {
69 return RunnerFunc(func(ctx context.Context, sink Sink) error { return nil })
70 })
71 a, err := f.Spawn(context.Background(), "share me")
72 if err != nil {
73 t.Fatal(err)
74 }
75 a.Wait()
76 if a.Dir != f.cfg.RepoRoot {
77 t.Fatalf("Dir = %q; want RepoRoot %q", a.Dir, f.cfg.RepoRoot)
78 }
79}
80
81func TestSpawnEmptyTaskFails(t *testing.T) {
82 f := newTestSwarm(t, func(a *Agent) Runner {
83 return RunnerFunc(func(ctx context.Context, sink Sink) error { return nil })
84 })
85 if _, err := f.Spawn(context.Background(), " "); err == nil {
86 t.Fatal("expected error on empty task")
87 }
88}
89
90func TestRunnerErrorMarksFailed(t *testing.T) {
91 wantErr := errors.New("boom")
92 f := newTestSwarm(t, func(a *Agent) Runner {
93 return RunnerFunc(func(ctx context.Context, sink Sink) error { return wantErr })
94 })
95 a, _ := f.Spawn(context.Background(), "explode")
96 a.Wait()
97 if a.Status() != StatusFailed {
98 t.Fatalf("status %s; want failed", a.Status())
99 }
100 if !errors.Is(a.Err(), wantErr) {
101 t.Fatalf("err = %v", a.Err())
102 }
103 if !strings.Contains(a.Activity(), "boom") {
104 t.Fatalf("activity = %q; want it to mention the error", a.Activity())
105 }
106}
107
108func TestStopCancelsRunningAgent(t *testing.T) {
109 started := make(chan struct{})
110 f := newTestSwarm(t, func(a *Agent) Runner {
111 return RunnerFunc(func(ctx context.Context, sink Sink) error {
112 close(started)
113 <-ctx.Done()
114 return ctx.Err()
115 })
116 })
117 a, _ := f.Spawn(context.Background(), "long")
118 <-started
119 if err := f.Stop(a.ID); err != nil {
120 t.Fatal(err)
121 }
122 a.Wait()
123 if a.Status() != StatusKilled {
124 t.Fatalf("status = %s; want killed", a.Status())
125 }
126}
127
128func TestStopAfterDoneIsNoop(t *testing.T) {
129 f := newTestSwarm(t, func(a *Agent) Runner {
130 return RunnerFunc(func(ctx context.Context, sink Sink) error { return nil })
131 })
132 a, _ := f.Spawn(context.Background(), "quick")
133 a.Wait()
134 if err := f.Stop(a.ID); err != nil {
135 t.Fatalf("stop after done: %v", err)
136 }
137 if a.Status() != StatusDone {
138 t.Fatalf("status flipped to %s", a.Status())
139 }
140}
141
142func TestGetPrefixMatch(t *testing.T) {
143 f := newTestSwarm(t, func(a *Agent) Runner {
144 return RunnerFunc(func(ctx context.Context, sink Sink) error { return nil })
145 })
146 a, _ := f.Spawn(context.Background(), "alpha task")
147 a.Wait()
148 // Full id works.
149 if got := f.Get(a.ID); got != a {
150 t.Fatal("get by full id failed")
151 }
152 // Slug prefix works as long as it's unique.
153 if got := f.Get("alpha"); got != a {
154 t.Fatal("get by prefix failed")
155 }
156 // Bogus id returns nil.
157 if got := f.Get("zzz-nope"); got != nil {
158 t.Fatalf("expected nil; got %#v", got)
159 }
160}
161
162func TestRemoveRequiresTerminalState(t *testing.T) {
163 hold := make(chan struct{})
164 f := newTestSwarm(t, func(a *Agent) Runner {
165 return RunnerFunc(func(ctx context.Context, sink Sink) error {
166 <-hold
167 return nil
168 })
169 })
170 a, _ := f.Spawn(context.Background(), "still going")
171 // Wait for run goroutine to flip to running.
172 for i := 0; i < 100 && a.Status() != StatusRunning; i++ {
173 time.Sleep(time.Millisecond)
174 }
175 if err := f.Remove(a.ID); err == nil {
176 t.Fatal("remove of running agent should fail")
177 }
178 close(hold)
179 a.Wait()
180 if err := f.Remove(a.ID); err != nil {
181 t.Fatalf("remove after done: %v", err)
182 }
183 if got := f.Get(a.ID); got != nil {
184 t.Fatal("agent still present after remove")
185 }
186}
187
188func TestSnapshotIsStableAcrossAccess(t *testing.T) {
189 var wg sync.WaitGroup
190 wg.Add(1)
191 f := newTestSwarm(t, func(a *Agent) Runner {
192 return RunnerFunc(func(ctx context.Context, sink Sink) error {
193 defer wg.Done()
194 for i := 0; i < 200; i++ {
195 sink.Transcript("chunk")
196 sink.Activity("step")
197 }
198 return nil
199 })
200 })
201 a, _ := f.Spawn(context.Background(), "race")
202 // Hammer Snapshot while the runner is writing; the -race detector
203 // is the real assertion here.
204 stop := make(chan struct{})
205 go func() {
206 for {
207 select {
208 case <-stop:
209 return
210 default:
211 _ = a.Snapshot()
212 }
213 }
214 }()
215 wg.Wait()
216 a.Wait()
217 close(stop)
218 if a.Status() != StatusDone {
219 t.Fatalf("status = %s", a.Status())
220 }
221}
222
223func TestTaskSlug(t *testing.T) {
224 cases := map[string]string{
225 "fix the login form": "fix-the-login-form",
226 " weird --- spaces!! ": "weird-spaces",
227 "": "agent",
228 "a-very-long-task-name-that-overflows": "a-very-long-task-name-th",
229 }
230 for in, want := range cases {
231 if got := taskSlug(in); got != want {
232 t.Errorf("taskSlug(%q) = %q; want %q", in, got, want)
233 }
234 }
235}
236
237func TestSnapshotAllSorted(t *testing.T) {
238 f := newTestSwarm(t, func(a *Agent) Runner {
239 return RunnerFunc(func(ctx context.Context, sink Sink) error { return nil })
240 })
241 a1, _ := f.Spawn(context.Background(), "first")
242 // Force second spawn into a later nanosecond bucket.
243 time.Sleep(2 * time.Millisecond)
244 a2, _ := f.Spawn(context.Background(), "second")
245 a1.Wait()
246 a2.Wait()
247 snaps := f.SnapshotAll()
248 if len(snaps) != 2 {
249 t.Fatalf("want 2 snapshots; got %d", len(snaps))
250 }
251 if !snaps[0].Started.Before(snaps[1].Started) && !snaps[0].Started.Equal(snaps[1].Started) {
252 t.Fatal("snapshots not in spawn order")
253 }
254}