Skip to content

refactor(testutil): extract shared WaitFor helper, deduplicate test poll-until helpers #1037

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yunaremaia wants to merge 2 commits into
base: main
Choose a base branch
from
Open

refactor(testutil): extract shared WaitFor helper, deduplicate test poll-until helpers #1037

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ceb231
refactor(testutil): extract shared WaitFor helper, deduplicate test p…
yunaremaia Sep 8, 2026
f57dc52
fix: add descriptive wait messages to daemon tests
yunaremaia Sep 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 4 additions & 14 deletions internal/daemon/pool_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ import (
"sync/atomic"
"testing"
"time"

"github.com/Gitlawb/zero/internal/testutil"
)

// --- test doubles ---------------------------------------------------------
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down
6 changes: 4 additions & 2 deletions internal/daemon/session_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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, "", 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")
}
Expand Down
48 changes: 19 additions & 29 deletions internal/swarm/lifecycle_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand All @@ -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
})
Expand All @@ -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)
Expand Down Expand Up @@ -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())
}
Expand All @@ -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
})
Expand All @@ -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
})
Expand All @@ -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
})
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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
})
Expand All @@ -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
Expand All @@ -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
})
Expand All @@ -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
})
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
18 changes: 10 additions & 8 deletions internal/swarm/scheduler_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
Expand All @@ -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
})
Expand All @@ -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
})
Expand Down Expand Up @@ -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)
Expand Down
13 changes: 7 additions & 6 deletions internal/swarm/tools_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"time"
"unicode/utf8"

"github.com/Gitlawb/zero/internal/testutil"
"github.com/Gitlawb/zero/internal/tools"
)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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
})
Expand Down Expand Up @@ -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 {
Expand All @@ -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
})
Expand Down
Loading