mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
test: add testutil.WaitBuffer and replace time.Sleep in tests (#22922)
WaitBuffer is a thread-safe io.Writer that supports blocking until accumulated output matches a substring or custom predicate. It replaces ad-hoc safeBuffer/syncWriter types and time.Sleep-based poll loops in tests with signal-driven waits. - WaitFor/WaitForNth/WaitForCond for blocking on output - Replace custom buffer types in cli/sync_test.go and provisionersdk/agent_test.go - Convert time.Sleep poll loops to require.Eventually/require.Never in cli/ssh_test.go, coderd/activitybump_test.go, coderd/workspaceagentsrpc_test.go, workspaceproxy_test.go, and scaletest tests
This commit is contained in:
@@ -0,0 +1,117 @@
|
||||
package testutil
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// WaitBuffer is a thread-safe buffer (io.Writer) that supports
|
||||
// blocking until the accumulated content matches a condition.
|
||||
// It is intended for tests that need to wait for specific output
|
||||
// from a command or process before proceeding.
|
||||
//
|
||||
// WaitBuffer is safe for concurrent use. Multiple goroutines may
|
||||
// write to it, and WaitFor/WaitForCond may be called from any
|
||||
// goroutine.
|
||||
type WaitBuffer struct {
|
||||
mu sync.Mutex
|
||||
buf bytes.Buffer
|
||||
waiters []*wbWaiter
|
||||
}
|
||||
|
||||
type wbWaiter struct {
|
||||
cond func(string) bool
|
||||
ch chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// NewWaitBuffer returns a new WaitBuffer. It can be used as a
|
||||
// plain thread-safe io.Writer even if WaitFor is never called.
|
||||
func NewWaitBuffer() *WaitBuffer {
|
||||
return &WaitBuffer{}
|
||||
}
|
||||
|
||||
// Write implements io.Writer. It is safe for concurrent use.
|
||||
func (wb *WaitBuffer) Write(p []byte) (int, error) {
|
||||
wb.mu.Lock()
|
||||
defer wb.mu.Unlock()
|
||||
|
||||
n, err := wb.buf.Write(p)
|
||||
s := wb.buf.String()
|
||||
for _, w := range wb.waiters {
|
||||
if w.cond(s) {
|
||||
w.once.Do(func() { close(w.ch) })
|
||||
}
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
// WaitFor blocks until the accumulated output contains signal or
|
||||
// ctx expires. Returns nil on match, ctx.Err() on timeout.
|
||||
// Safe to call from any goroutine.
|
||||
func (wb *WaitBuffer) WaitFor(ctx context.Context, signal string) error {
|
||||
return wb.WaitForNth(ctx, signal, 1)
|
||||
}
|
||||
|
||||
// WaitForNth blocks until the accumulated output contains at least
|
||||
// n occurrences of signal, or ctx expires. Returns nil on match,
|
||||
// ctx.Err() on timeout. Safe to call from any goroutine.
|
||||
func (wb *WaitBuffer) WaitForNth(ctx context.Context, signal string, n int) error {
|
||||
return wb.WaitForCond(ctx, func(s string) bool {
|
||||
return strings.Count(s, signal) >= n
|
||||
})
|
||||
}
|
||||
|
||||
// WaitForCond blocks until cond returns true for the accumulated
|
||||
// output, or ctx expires. Returns nil on match, ctx.Err() on
|
||||
// timeout. Safe to call from any goroutine.
|
||||
func (wb *WaitBuffer) WaitForCond(ctx context.Context, cond func(string) bool) error {
|
||||
wb.mu.Lock()
|
||||
if cond(wb.buf.String()) {
|
||||
wb.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
w := &wbWaiter{
|
||||
cond: cond,
|
||||
ch: make(chan struct{}),
|
||||
}
|
||||
wb.waiters = append(wb.waiters, w)
|
||||
wb.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-w.ch:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// RequireWaitFor blocks until the accumulated output contains
|
||||
// signal or ctx expires. On timeout, fails the test with a
|
||||
// message showing what was expected and what was written so far.
|
||||
//
|
||||
// Safety: Must only be called from the Go routine that created
|
||||
// `t`.
|
||||
func (wb *WaitBuffer) RequireWaitFor(ctx context.Context, t testing.TB, signal string) {
|
||||
t.Helper()
|
||||
if err := wb.WaitFor(ctx, signal); err != nil {
|
||||
t.Fatalf("WaitBuffer: signal %q not found; buffer contents:\n%s", signal, wb.String())
|
||||
}
|
||||
}
|
||||
|
||||
// Bytes returns a copy of the accumulated output.
|
||||
func (wb *WaitBuffer) Bytes() []byte {
|
||||
wb.mu.Lock()
|
||||
defer wb.mu.Unlock()
|
||||
return bytes.Clone(wb.buf.Bytes())
|
||||
}
|
||||
|
||||
// String returns the accumulated output as a string.
|
||||
func (wb *WaitBuffer) String() string {
|
||||
wb.mu.Lock()
|
||||
defer wb.mu.Unlock()
|
||||
return wb.buf.String()
|
||||
}
|
||||
Reference in New Issue
Block a user