diff --git a/internal/daemon/pool_test.go b/internal/daemon/pool_test.go index 7ced2bef4..a2138495e 100644 --- a/internal/daemon/pool_test.go +++ b/internal/daemon/pool_test.go @@ -7,6 +7,8 @@ import ( "sync/atomic" "testing" "time" + + "github.com/Gitlawb/zero/internal/testutil" ) // --- test doubles --------------------------------------------------------- @@ -191,7 +193,7 @@ func TestPoolQueuesWhenFull(t *testing.T) { }() <-started // Wait until the first run holds the only slot. - waitFor(t, func() bool { return pool.QueueDepth() == 1 }) + testutil.WaitFor(t, "queue depth reaches 1", func() bool { return pool.QueueDepth() == 1 }) secondDone := make(chan struct{}) go func() { @@ -224,7 +226,7 @@ func TestPoolDrainKillsStraggler(t *testing.T) { _, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{}) close(runDone) }() - waitFor(t, func() bool { return pool.QueueDepth() == 1 }) + testutil.WaitFor(t, "queue depth reaches 1", func() bool { return pool.QueueDepth() == 1 }) pool.Drain() // KillTimeout elapses, straggler is force-killed if atomic.LoadInt32(&straggler.killed) != 1 { @@ -242,18 +244,6 @@ func TestPoolDrainKillsStraggler(t *testing.T) { } } -func waitFor(t *testing.T, cond func() bool) { - t.Helper() - deadline := time.Now().Add(2 * time.Second) - for time.Now().Before(deadline) { - if cond() { - return - } - time.Sleep(time.Millisecond) - } - t.Fatal("condition not met within timeout") -} - func TestPoolRunSurfacesStdoutReadError(t *testing.T) { // A worker stdout read error must surface as a failure (not be swallowed and // reported as clean success), and the worker is killed so it can't block on an diff --git a/internal/daemon/session_test.go b/internal/daemon/session_test.go index c038638f2..293016980 100644 --- a/internal/daemon/session_test.go +++ b/internal/daemon/session_test.go @@ -6,6 +6,8 @@ import ( "sync/atomic" "testing" "time" + + "github.com/Gitlawb/zero/internal/testutil" ) func drain(t *testing.T, buffered []string, live <-chan string) []string { @@ -111,7 +113,7 @@ func TestSessionLeaseQueuesWhenPoolFull(t *testing.T) { mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool}) sa, _ := mgr.Start(context.Background(), WorkerSpec{Session: "a"}) - waitFor(t, func() bool { return sa.State() == SessionRunning }) + testutil.WaitFor(t, "session reaches running state", func() bool { return sa.State() == SessionRunning }) sb, _ := mgr.Start(context.Background(), WorkerSpec{Session: "b"}) // b must stay queued (its worker not yet launched) while a holds the slot. @@ -197,7 +199,7 @@ func TestSessionManagerKeepsRunningOverCap(t *testing.T) { s1, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r1"}) s2, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r2"}) - waitFor(t, func() bool { return s1.State() == SessionRunning && s2.State() == SessionRunning }) + testutil.WaitFor(t, "sessions running", func() bool { return s1.State() == SessionRunning && s2.State() == SessionRunning }) if _, ok := mgr.Get("r1"); !ok { t.Fatal("a running session must never be evicted, even past the cap") } diff --git a/internal/swarm/lifecycle_test.go b/internal/swarm/lifecycle_test.go index e119550c1..4822e4a59 100644 --- a/internal/swarm/lifecycle_test.go +++ b/internal/swarm/lifecycle_test.go @@ -7,6 +7,8 @@ import ( "sync" "testing" "time" + + "github.com/Gitlawb/zero/internal/testutil" ) // controllableLauncher records every launched spec and lets a test control each @@ -80,18 +82,6 @@ func newSwarmFor(t *testing.T, l MemberLauncher) *Swarm { return sw } -func waitFor(t *testing.T, what string, cond func() bool) { - t.Helper() - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if cond() { - return - } - time.Sleep(5 * time.Millisecond) - } - t.Fatalf("timed out waiting for %s", what) -} - func okFor(spec MemberSpec, _ int) (MemberResult, error) { return MemberResult{Result: "ok:" + spec.Task, SessionID: "sess-" + spec.ID}, nil } @@ -103,7 +93,7 @@ func TestSpawnCompletes(t *testing.T) { if err != nil { t.Fatalf("Spawn: %v", err) } - waitFor(t, "task done", func() bool { + testutil.WaitFor(t, "task done", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusDone }) @@ -120,7 +110,7 @@ func TestSpawnInheritsPolicy(t *testing.T) { if err != nil { t.Fatalf("Spawn: %v", err) } - waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 }) + testutil.WaitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 }) spec := l.recorded()[0] if spec.Model != "orch-model" { t.Fatalf("member model = %q, want inherited orch-model", spec.Model) @@ -153,7 +143,7 @@ func TestConcurrencyCapAndQueueDrains(t *testing.T) { } // Release everyone; the queue should drain one-per-slot until all are done. close(gate) - waitFor(t, "all tasks done", func() bool { return sw.Coordinator().Summarize().Done == 5 }) + testutil.WaitFor(t, "all tasks done", func() bool { return sw.Coordinator().Summarize().Done == 5 }) if team.Running() != 0 || team.QueueDepth() != 0 { t.Fatalf("after drain running=%d queue=%d, want 0/0", team.Running(), team.QueueDepth()) } @@ -171,7 +161,7 @@ func TestRetryOnTemporaryError(t *testing.T) { }) sw := newSwarmFor(t, l) id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "") - waitFor(t, "task recovered", func() bool { + testutil.WaitFor(t, "task recovered", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusDone }) @@ -190,7 +180,7 @@ func TestRetryExhaustionFails(t *testing.T) { }) sw := newSwarmFor(t, l) id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "") - waitFor(t, "task failed", func() bool { + testutil.WaitFor(t, "task failed", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusFailed }) @@ -205,7 +195,7 @@ func TestPermanentErrorNoRetry(t *testing.T) { }) sw := newSwarmFor(t, l) id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "") - waitFor(t, "task failed", func() bool { + testutil.WaitFor(t, "task failed", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusFailed }) @@ -291,7 +281,7 @@ func TestClosePreventsMemberRetry(t *testing.T) { sw.Close() close(closed) }() - waitFor(t, "swarm closed state", func() bool { + testutil.WaitFor(t, "swarm closed state", func() bool { sw.lifecycleMu.RLock() defer sw.lifecycleMu.RUnlock() return sw.closed @@ -320,7 +310,7 @@ func TestCloseWaitsForMemberWatchers(t *testing.T) { if err != nil { t.Fatalf("Spawn: %v", err) } - waitFor(t, "task running", func() bool { + testutil.WaitFor(t, "task running", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusRunning }) @@ -360,7 +350,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) { sw := newSwarmFor(t, l) pol := Policy{Model: "m"} origID, _ := sw.Spawn(pol, "team", "teammate", "original task", "/w") - waitFor(t, "original running", func() bool { + testutil.WaitFor(t, "original running", func() bool { task, ok := sw.Coordinator().Get(origID) return ok && task.Status == StatusRunning }) @@ -384,7 +374,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) { } // The new member carries the handoff note in its task and preserves cwd. close(gate) - waitFor(t, "spec for new member", func() bool { + testutil.WaitFor(t, "spec for new member", func() bool { for _, s := range l.recorded() { if s.ID == newID { return true @@ -400,7 +390,7 @@ func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) { } } // A handoff of an already-terminal task is rejected. - waitFor(t, "new task done", func() bool { + testutil.WaitFor(t, "new task done", func() bool { task, ok := sw.Coordinator().Get(newID) return ok && task.Status == StatusDone }) @@ -427,7 +417,7 @@ func TestAdoptOrphans(t *testing.T) { t.Fatalf("adopted = %v, want [orphan-1]", adopted) } // The orphan is relaunched under a fresh agent and completes. - waitFor(t, "orphan completed", func() bool { + testutil.WaitFor(t, "orphan completed", func() bool { task, ok := sw.Coordinator().Get("orphan-1") return ok && task.Status == StatusDone }) @@ -468,7 +458,7 @@ func TestCollectScopesToTeam(t *testing.T) { sw := newSwarmFor(t, l) a, _ := sw.Spawn(Policy{}, "alpha", "teammate", "ta", "") _, _ = sw.Spawn(Policy{}, "beta", "teammate", "tb", "") - waitFor(t, "alpha done", func() bool { + testutil.WaitFor(t, "alpha done", func() bool { task, ok := sw.Coordinator().Get(a) return ok && task.Status == StatusDone }) @@ -614,7 +604,7 @@ func TestAdmittedDispatchDoesNotLaunchAfterClose(t *testing.T) { sw.Close() close(closed) }() - waitFor(t, "swarm closed state", func() bool { + testutil.WaitFor(t, "swarm closed state", func() bool { sw.lifecycleMu.RLock() defer sw.lifecycleMu.RUnlock() return sw.closed @@ -674,7 +664,7 @@ func TestCloseBetweenLaunchPrecheckAndReturn(t *testing.T) { sw.Close() close(closed) }() - waitFor(t, "swarm closed state", func() bool { + testutil.WaitFor(t, "swarm closed state", func() bool { sw.lifecycleMu.RLock() defer sw.lifecycleMu.RUnlock() return sw.closed @@ -772,7 +762,7 @@ func TestCloseBetweenRetryLaunchPrecheckAndReturn(t *testing.T) { sw.Close() close(closed) }() - waitFor(t, "swarm closed state", func() bool { + testutil.WaitFor(t, "swarm closed state", func() bool { sw.lifecycleMu.RLock() defer sw.lifecycleMu.RUnlock() return sw.closed @@ -964,7 +954,7 @@ func TestCloseFailsQueuedSpecOnLateCreatedTeam(t *testing.T) { sw.Close() close(closed) }() - waitFor(t, "swarm closed state", func() bool { + testutil.WaitFor(t, "swarm closed state", func() bool { sw.lifecycleMu.RLock() defer sw.lifecycleMu.RUnlock() return sw.closed diff --git a/internal/swarm/scheduler_test.go b/internal/swarm/scheduler_test.go index d3b04be0b..91a6430af 100644 --- a/internal/swarm/scheduler_test.go +++ b/internal/swarm/scheduler_test.go @@ -4,6 +4,8 @@ import ( "math" "testing" "time" + + "github.com/Gitlawb/zero/internal/testutil" ) // testTicker returns a ticker factory backed by a single unbounded-handshake @@ -39,11 +41,11 @@ func TestSchedulerFiresAndCountsRuns(t *testing.T) { for i := 0; i < 3; i++ { ticks <- time.Time{} want := i + 1 - waitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == want }) + testutil.WaitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == want }) } // After MaxRuns the job retires itself. - waitFor(t, "job retired", func() bool { _, ok := findJob(sched.List(), id); return !ok }) + testutil.WaitFor(t, "job retired", func() bool { _, ok := findJob(sched.List(), id); return !ok }) if got := len(l.recorded()); got != 3 { t.Fatalf("spawned %d members, want 3", got) } @@ -68,11 +70,11 @@ func TestSchedulerSkipsWhilePreviousRuns(t *testing.T) { // Fire 1: spawns and the member stays running (gated). ticks <- time.Time{} - waitFor(t, "first spawn", func() bool { return len(l.recorded()) == 1 }) + testutil.WaitFor(t, "first spawn", func() bool { return len(l.recorded()) == 1 }) // Fire 2: previous still running => skipped, no new spawn. ticks <- time.Time{} - waitFor(t, "skip recorded", func() bool { + testutil.WaitFor(t, "skip recorded", func() bool { j, ok := findJob(sched.List(), id) return ok && j.Skipped == 1 }) @@ -82,14 +84,14 @@ func TestSchedulerSkipsWhilePreviousRuns(t *testing.T) { // Release the first member, then fire 3: previous terminal => spawns again. close(gate) - waitFor(t, "first done", func() bool { return sw.Coordinator().Summarize().Done == 1 }) + testutil.WaitFor(t, "first done", func() bool { return sw.Coordinator().Summarize().Done == 1 }) ticks <- time.Time{} - waitFor(t, "second spawn", func() bool { return len(l.recorded()) == 2 }) + testutil.WaitFor(t, "second spawn", func() bool { return len(l.recorded()) == 2 }) // fireIfIdle's spawn (what "second spawn" observes via the launcher) and // run's subsequent job.incRuns() are sequential but distinct steps in the // scheduler's goroutine; wait for Runs itself rather than assuming the // launcher recording it means the job's counter is updated too. - waitFor(t, "second run recorded", func() bool { + testutil.WaitFor(t, "second run recorded", func() bool { j, ok := findJob(sched.List(), id) return ok && j.Runs == 2 }) @@ -215,7 +217,7 @@ func TestSchedulerDailyRecomputesNextDelay(t *testing.T) { } for i := 0; i < 2; i++ { ticks <- time.Time{} - waitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == i+1 }) + testutil.WaitFor(t, "task completed", func() bool { return sw.Coordinator().Summarize().Done == i+1 }) } if len(delays) < 2 { t.Fatalf("expected at least two requested delays, got %v", delays) diff --git a/internal/swarm/tools_test.go b/internal/swarm/tools_test.go index 82dccaa75..b8dd99a55 100644 --- a/internal/swarm/tools_test.go +++ b/internal/swarm/tools_test.go @@ -8,6 +8,7 @@ import ( "time" "unicode/utf8" + "github.com/Gitlawb/zero/internal/testutil" "github.com/Gitlawb/zero/internal/tools" ) @@ -94,7 +95,7 @@ func TestSpawnToolThroughRegistry(t *testing.T) { if id == "" { t.Fatal("spawn must return a task_id in Meta") } - waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 }) + testutil.WaitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 }) spec := l.recorded()[0] if spec.Model != "m1" || spec.Cwd != "/work" { t.Fatalf("policy/cwd not threaded into member: model=%q cwd=%q", spec.Model, spec.Cwd) @@ -165,7 +166,7 @@ func TestStatusAndCollectTools(t *testing.T) { "agent_type": "teammate", "task": "compute", "team": "alpha", }, tools.RunOptions{PermissionGranted: true, Model: "m"}) id := spawn.Meta["task_id"] - waitFor(t, "task done", func() bool { + testutil.WaitFor(t, "task done", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusDone }) @@ -213,7 +214,7 @@ func TestCollectBlocksUntilMembersFinish(t *testing.T) { "agent_type": "teammate", "task": "compute", "team": "alpha", }, tools.RunOptions{PermissionGranted: true, Model: "m"}) id := spawn.Meta["task_id"] - waitFor(t, "member running", func() bool { + testutil.WaitFor(t, "member running", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusRunning }) @@ -258,7 +259,7 @@ func TestCollectReturnsPartialOnTimeout(t *testing.T) { "agent_type": "teammate", "task": "compute", "team": "alpha", }, tools.RunOptions{PermissionGranted: true, Model: "m"}) id := spawn.Meta["task_id"] - waitFor(t, "member running", func() bool { + testutil.WaitFor(t, "member running", func() bool { task, ok := sw.Coordinator().Get(id) return ok && task.Status == StatusRunning }) @@ -324,7 +325,7 @@ func TestNewWithUserDefinitions(t *testing.T) { if _, err := sw.Spawn(Policy{Model: "m"}, "team", "researcher", "find things", ""); err != nil { t.Fatalf("Spawn user-defined agent: %v", err) } - waitFor(t, "researcher used", func() bool { return used }) + testutil.WaitFor(t, "researcher used", func() bool { return used }) // A bad definition makes New fail closed. if _, err := New(Options{BaseDir: t.TempDir(), Launcher: newLauncher(okFor), Definitions: []Definition{{AgentType: " "}}}); err == nil { @@ -341,7 +342,7 @@ func TestHandoffToolThroughRegistry(t *testing.T) { "agent_type": "teammate", "task": "long", "team": "alpha", }, tools.RunOptions{PermissionGranted: true, Model: "m"}) origID := spawn.Meta["task_id"] - waitFor(t, "running", func() bool { + testutil.WaitFor(t, "running", func() bool { task, ok := sw.Coordinator().Get(origID) return ok && task.Status == StatusRunning }) diff --git a/internal/testutil/testutil.go b/internal/testutil/testutil.go new file mode 100644 index 000000000..9fb56e2be --- /dev/null +++ b/internal/testutil/testutil.go @@ -0,0 +1,22 @@ +// Package testutil provides shared test helpers for Zero. +package testutil + +import ( + "testing" + "time" +) + +// WaitFor polls cond until it returns true or the default deadline (5s) is +// reached. The what argument names what is being waited for in the failure +// message. Deadline and interval are generous enough for loaded CI runners. +func WaitFor(t *testing.T, what string, cond func() bool) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if cond() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("timed out waiting for %s", what) +}