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

cli/ssh_test.go

@@ -180,15 +180,11 @@ func TestSSH(t *testing.T) {
 
 		// Delay until workspace is starting, otherwise the agent may be
 		// booted due to outdated build.
-		var err error
-		for {
+		require.Eventually(t, func() bool {
+			var err error
 			workspace, err = client.Workspace(ctx, workspace.ID)
-			require.NoError(t, err)
-			if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart {
-				break
-			}
-			time.Sleep(testutil.IntervalFast)
-		}
+			return err == nil && workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart
+		}, testutil.WaitShort, testutil.IntervalFast)
 
 		// When the agent connects, the workspace was started, and we should
 		// have access to the shell.
@@ -763,15 +759,11 @@ func TestSSH(t *testing.T) {
 
 		// Delay until workspace is starting, otherwise the agent may be
 		// booted due to outdated build.
-		var err error
-		for {
+		require.Eventually(t, func() bool {
+			var err error
 			workspace, err = client.Workspace(ctx, workspace.ID)
-			require.NoError(t, err)
-			if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart {
-				break
-			}
-			time.Sleep(testutil.IntervalFast)
-		}
+			return err == nil && workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionStart
+		}, testutil.WaitShort, testutil.IntervalFast)
 
 		// When the agent connects, the workspace was started, and we should
 		// have access to the shell.

cli/sync_test.go

@@ -6,8 +6,6 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
-	"strings"
-	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -104,20 +102,11 @@ func TestSyncCommands_Golden(t *testing.T) {
 		require.NoError(t, err)
 		client.Close()
 
-		// Use a writer that signals when the "Waiting" message has been
-		// written, so the goroutine can complete the dependency at the
-		// right time without relying on time.Sleep.
-		outBuf := newSyncWriter("Waiting")
-
-		// Start a goroutine to complete the dependency once the start
-		// command has printed its waiting message.
+		outBuf := testutil.NewWaitBuffer()
 		done := make(chan error, 1)
 		go func() {
-			// Block until the command prints the waiting message.
-			select {
-			case <-outBuf.matched:
-			case <-ctx.Done():
-				done <- ctx.Err()
+			if err := outBuf.WaitFor(ctx, "Waiting"); err != nil {
+				done <- err
 				return
 			}
 
@@ -339,36 +328,3 @@ func TestSyncCommands_Golden(t *testing.T) {
 		clitest.TestGoldenFile(t, "TestSyncCommands_Golden/status_json_format", outBuf.Bytes(), nil)
 	})
 }
-
-// syncWriter is a thread-safe io.Writer that wraps a bytes.Buffer and
-// closes a channel when the written content contains a signal string.
-type syncWriter struct {
-	mu        sync.Mutex
-	buf       bytes.Buffer
-	signal    string
-	matched   chan struct{}
-	closeOnce sync.Once
-}
-
-func newSyncWriter(signal string) *syncWriter {
-	return &syncWriter{
-		signal:  signal,
-		matched: make(chan struct{}),
-	}
-}
-
-func (w *syncWriter) Write(p []byte) (int, error) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	n, err := w.buf.Write(p)
-	if w.signal != "" && strings.Contains(w.buf.String(), w.signal) {
-		w.closeOnce.Do(func() { close(w.matched) })
-	}
-	return n, err
-}
-
-func (w *syncWriter) Bytes() []byte {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	return w.buf.Bytes()
-}

coderd/activitybump_test.go

@@ -116,10 +116,10 @@ func TestWorkspaceActivityBump(t *testing.T) {
 				// is required. The Activity Bump behavior is also coupled with
 				// Last Used, so it would be obvious to the user if we
 				// are falsely recognizing activity.
-				time.Sleep(testutil.IntervalMedium)
-				workspace, err = client.Workspace(ctx, workspace.ID)
-				require.NoError(t, err)
-				require.Equal(t, workspace.LatestBuild.Deadline.Time, firstDeadline)
+				require.Never(t, func() bool {
+					workspace, err = client.Workspace(ctx, workspace.ID)
+					return err == nil && !workspace.LatestBuild.Deadline.Time.Equal(firstDeadline)
+				}, testutil.IntervalMedium, testutil.IntervalFast, "deadline should not change")
 				return
 			}
 

coderd/workspaceagentsrpc_test.go

@@ -233,17 +233,16 @@ func TestWorkspaceAgentRPCRole(t *testing.T) {
 
 		// Close the connection and give the server time to process.
 		_ = conn.Close()
-		time.Sleep(100 * time.Millisecond)
 
-		// Verify that connectivity timestamps were never set.
-		agent, err := db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), r.Agents[0].ID)
-		require.NoError(t, err)
-		assert.False(t, agent.FirstConnectedAt.Valid,
-			"first_connected_at should NOT be set for non-agent role")
-		assert.False(t, agent.LastConnectedAt.Valid,
-			"last_connected_at should NOT be set for non-agent role")
-		assert.False(t, agent.DisconnectedAt.Valid,
-			"disconnected_at should NOT be set for non-agent role")
+		// Verify that connectivity timestamps were never set
+		// (first_connected_at, last_connected_at, disconnected_at).
+		require.Never(t, func() bool {
+			agent, err := db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), r.Agents[0].ID)
+			if err != nil {
+				return false
+			}
+			return agent.FirstConnectedAt.Valid || agent.LastConnectedAt.Valid || agent.DisconnectedAt.Valid
+		}, testutil.IntervalMedium, testutil.IntervalFast, "connectivity timestamps should NOT be set for non-agent role")
 	})
 
 	// NOTE: Backward compatibility (empty role) is implicitly tested by

docs/about/contributing/backend.md

@@ -118,6 +118,7 @@ The Coder backend includes a rich suite of unit and end-to-end tests. A variety
   * [port.go](https://github.com/coder/coder/blob/main/testutil/port.go): select a free random port
   * [prometheus.go](https://github.com/coder/coder/blob/main/testutil/prometheus.go): validate Prometheus metrics with expected values
   * [pty.go](https://github.com/coder/coder/blob/main/testutil/pty.go): read output from a terminal until a condition is met
+  * [wait_buffer.go](https://github.com/coder/coder/blob/main/testutil/wait_buffer.go): thread-safe `io.Writer` that blocks until accumulated output contains a signal (`WaitFor`, `WaitForNth`, `WaitForCond`)
 
 ### [dbtestutil](https://github.com/coder/coder/tree/main/coderd/database/dbtestutil)
 

enterprise/coderd/workspaceproxy_test.go

@@ -586,9 +586,13 @@ func TestProxyRegisterDeregister(t *testing.T) {
 		proxyClient := wsproxysdk.New(client.URL, createRes.ProxyToken)
 
 		for i := 0; i < 100; i++ {
-			ok := false
-			for j := 0; j < 2; j++ {
-				registerRes, err := proxyClient.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
+			// Sibling replica count may not be immediately consistent.
+			// In production, proxies re-register every 30s and
+			// Kubernetes rolls out gradually, so this is benign.
+			var registerRes wsproxysdk.RegisterWorkspaceProxyResponse
+			require.Eventually(t, func() bool {
+				var err error
+				registerRes, err = proxyClient.RegisterWorkspaceProxy(ctx, wsproxysdk.RegisterWorkspaceProxyRequest{
 					AccessURL:           "https://proxy.coder.test",
 					WildcardHostname:    "*.proxy.coder.test",
 					DerpEnabled:         true,
@@ -598,25 +602,11 @@ func TestProxyRegisterDeregister(t *testing.T) {
 					ReplicaRelayAddress: fmt.Sprintf("http://127.0.0.1:%d", 8080+i),
 					Version:             buildinfo.Version(),
 				})
-				require.NoErrorf(t, err, "register proxy %d", i)
-
-				// If the sibling replica count is wrong, try again. The impact
-				// of this not being immediate is that proxies may not function
-				// as DERP relays until they register again in 30 seconds.
-				//
-				// In the real world, replicas will not be registering this
-				// quickly. Kubernetes rolls out gradually in practice.
-				if len(registerRes.SiblingReplicas) != i {
-					t.Logf("%d: expected %d siblings, got %d", i, i, len(registerRes.SiblingReplicas))
-					time.Sleep(100 * time.Millisecond)
-					continue
+				if err != nil {
+					return false
 				}
-
-				ok = true
-				break
-			}
-
-			require.True(t, ok, "expected to register replica %d", i)
+				return len(registerRes.SiblingReplicas) == i
+			}, testutil.WaitShort, testutil.IntervalMedium, "expected to register replica %d with %d siblings", i, i)
 		}
 	})
 

provisionersdk/agent_test.go

@@ -7,7 +7,6 @@
 package provisionersdk_test
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"net/http"
@@ -48,13 +47,13 @@ func TestAgentScript(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitShort)
 		script := serveScript(t, bashEcho)
 
-		var output safeBuffer
+		output := testutil.NewWaitBuffer()
 		// This is intentionally ran in single quotes to mimic how a customer may
 		// embed our script. Our scripts should not include any single quotes.
 		// nolint:gosec
 		cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
-		cmd.Stdout = &output
-		cmd.Stderr = &output
+		cmd.Stdout = output
+		cmd.Stderr = output
 		require.NoError(t, cmd.Start())
 
 		err := cmd.Wait()
@@ -83,14 +82,14 @@ func TestAgentScript(t *testing.T) {
 		ctx := testutil.Context(t, testutil.WaitShort)
 		script := serveScript(t, unexpectedEcho)
 
-		var output safeBuffer
+		output := testutil.NewWaitBuffer()
 		// This is intentionally ran in single quotes to mimic how a customer may
 		// embed our script. Our scripts should not include any single quotes.
 		// nolint:gosec
 		cmd := exec.CommandContext(ctx, "sh", "-c", "sh -c '"+script+"'")
 		cmd.WaitDelay = time.Second
-		cmd.Stdout = &output
-		cmd.Stderr = &output
+		cmd.Stdout = output
+		cmd.Stderr = output
 		require.NoError(t, cmd.Start())
 
 		done := make(chan error, 1)
@@ -127,9 +126,7 @@ func TestAgentScript(t *testing.T) {
 
 		t.Log(output.String())
 
-		require.Eventually(t, func() bool {
-			return bytes.Contains(output.Bytes(), []byte("ERROR: Downloaded agent binary returned unexpected version output"))
-		}, testutil.WaitShort, testutil.IntervalSlow)
+		output.RequireWaitFor(ctx, t, "ERROR: Downloaded agent binary returned unexpected version output")
 	})
 }
 
@@ -155,33 +152,3 @@ func serveScript(t *testing.T, in string) string {
 	script = strings.ReplaceAll(script, "${AUTH_TYPE}", "token")
 	return script
 }
-
-// safeBuffer is a concurrency-safe bytes.Buffer
-type safeBuffer struct {
-	mu  sync.Mutex
-	buf bytes.Buffer
-}
-
-func (sb *safeBuffer) Write(p []byte) (n int, err error) {
-	sb.mu.Lock()
-	defer sb.mu.Unlock()
-	return sb.buf.Write(p)
-}
-
-func (sb *safeBuffer) Read(p []byte) (n int, err error) {
-	sb.mu.Lock()
-	defer sb.mu.Unlock()
-	return sb.buf.Read(p)
-}
-
-func (sb *safeBuffer) Bytes() []byte {
-	sb.mu.Lock()
-	defer sb.mu.Unlock()
-	return sb.buf.Bytes()
-}
-
-func (sb *safeBuffer) String() string {
-	sb.mu.Lock()
-	defer sb.mu.Unlock()
-	return sb.buf.String()
-}

scaletest/createworkspaces/run_test.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -541,19 +540,18 @@ func goEventuallyStartFakeAgent(ctx context.Context, t *testing.T, client *coder
 	go func() {
 		defer close(ch)
 		var workspace codersdk.Workspace
-		for {
+		if !assert.Eventually(t, func() bool {
 			res, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{})
-			if !assert.NoError(t, err) {
-				return
+			if err != nil {
+				return false
 			}
-			workspaces := res.Workspaces
-
-			if len(workspaces) == 1 {
-				workspace = workspaces[0]
-				break
+			if len(res.Workspaces) == 1 {
+				workspace = res.Workspaces[0]
+				return true
 			}
-
-			time.Sleep(testutil.IntervalMedium)
+			return false
+		}, testutil.WaitShort, testutil.IntervalMedium) {
+			return
 		}
 
 		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)

scaletest/workspacebuild/run_test.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -118,25 +117,23 @@ func Test_Runner(t *testing.T) {
 		// finish, then start the agents.
 		go func() {
 			var workspace codersdk.Workspace
-			for {
+			if !assert.Eventually(t, func() bool {
 				res, err := client.Workspaces(ctx, codersdk.WorkspaceFilter{
 					Owner: codersdk.Me,
 				})
-				if !assert.NoError(t, err) {
-					return
+				if err != nil {
+					return false
 				}
-				workspaces := res.Workspaces
-
-				if len(workspaces) == 1 {
-					workspace = workspaces[0]
-					break
+				if len(res.Workspaces) == 1 {
+					workspace = res.Workspaces[0]
+					return true
 				}
-
-				time.Sleep(100 * time.Millisecond)
+				return false
+			}, testutil.WaitShort, testutil.IntervalMedium) {
+				return
 			}
 
 			coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
-
 			// Start the three agents.
 			for i, authToken := range []string{authToken1, authToken2, authToken3} {
 				i := i + 1

testutil/wait_buffer.go

@@ -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()
+}

testutil/wait_buffer_test.go

@@ -0,0 +1,271 @@
+package testutil_test
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestWaitBuffer_WaitFor_Blocks(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = wb.WaitFor(ctx, "hello")
+	}()
+
+	// Write the signal after the goroutine is blocking.
+	_, err := wb.Write([]byte("hello"))
+	require.NoError(t, err)
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal("WaitFor did not unblock after signal was written")
+	}
+}
+
+func TestWaitBuffer_WaitFor_AlreadyPresent(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("already here"))
+	require.NoError(t, err)
+
+	// Signal is already in the buffer; WaitFor returns immediately.
+	require.NoError(t, wb.WaitFor(ctx, "already"))
+}
+
+func TestWaitBuffer_WaitFor_ContextExpired(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Already expired.
+
+	wb := testutil.NewWaitBuffer()
+	err := wb.WaitFor(ctx, "never")
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestWaitBuffer_WaitFor_MultipleWrites(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	// Write partial content that doesn't satisfy the condition.
+	_, err := wb.Write([]byte("hell"))
+	require.NoError(t, err)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = wb.WaitFor(ctx, "hello")
+	}()
+
+	// Complete the signal with a second write.
+	_, err = wb.Write([]byte("o"))
+	require.NoError(t, err)
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal("WaitFor did not unblock after multiple writes completed the signal")
+	}
+}
+
+func TestWaitBuffer_WaitForCond(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		// Wait until the buffer has at least 10 bytes.
+		_ = wb.WaitForCond(ctx, func(s string) bool {
+			return len(s) >= 10
+		})
+	}()
+
+	_, err := wb.Write([]byte("12345"))
+	require.NoError(t, err)
+	_, err = wb.Write([]byte("67890"))
+	require.NoError(t, err)
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal("WaitForCond did not unblock when condition was met")
+	}
+}
+
+func TestWaitBuffer_ConcurrentWrites(t *testing.T) {
+	t.Parallel()
+
+	wb := testutil.NewWaitBuffer()
+	var wg sync.WaitGroup
+	const writers = 10
+	const iterations = 100
+	wg.Add(writers)
+	for i := range writers {
+		go func() {
+			defer wg.Done()
+			for j := range iterations {
+				_, _ = wb.Write([]byte(fmt.Sprintf("w%d-%d ", i, j)))
+			}
+		}()
+	}
+	wg.Wait()
+
+	// Every write should have landed; verify no data was lost by
+	// checking the length is at least as large as expected.
+	assert.GreaterOrEqual(t, len(wb.Bytes()), writers*iterations)
+}
+
+func TestWaitBuffer_WaitFor_BackgroundGoroutine(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // Expire immediately.
+
+	wb := testutil.NewWaitBuffer()
+
+	// WaitFor from a background goroutine should return the
+	// context error rather than calling t.Fatal.
+	done := make(chan error, 1)
+	go func() {
+		done <- wb.WaitFor(ctx, "never")
+	}()
+
+	err := <-done
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestWaitBuffer_SequentialWaits(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+
+	_, err := wb.Write([]byte("first "))
+	require.NoError(t, err)
+	require.NoError(t, wb.WaitFor(ctx, "first"))
+
+	_, err = wb.Write([]byte("second"))
+	require.NoError(t, err)
+	require.NoError(t, wb.WaitFor(ctx, "second"))
+}
+
+func TestWaitBuffer_WaitForNth_Blocks(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("Foo "))
+	require.NoError(t, err)
+
+	// First occurrence is already present, but we want two.
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_ = wb.WaitForNth(ctx, "Foo", 2)
+	}()
+
+	_, err = wb.Write([]byte("Bar Foo"))
+	require.NoError(t, err)
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Fatal("WaitForNth did not unblock after second occurrence")
+	}
+}
+
+func TestWaitBuffer_WaitForNth_AlreadySatisfied(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("Foo Foo Foo"))
+	require.NoError(t, err)
+
+	// All three occurrences already present.
+	require.NoError(t, wb.WaitForNth(ctx, "Foo", 3))
+}
+
+func TestWaitBuffer_RequireWaitFor_Timeout(t *testing.T) {
+	t.Parallel()
+
+	// Use a mock testing.TB to capture the fatal call without
+	// killing the real test.
+	mock := &tbMock{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("some output"))
+	require.NoError(t, err)
+
+	wb.RequireWaitFor(ctx, mock, "missing-signal")
+	assert.True(t, mock.failed(), "expected RequireWaitFor to fail the mock test")
+}
+
+// tbMock is a minimal testing.TB that records Fatalf calls.
+type tbMock struct {
+	testing.TB // Embed to satisfy the interface.
+	mu         sync.Mutex
+	fatalCalls int
+}
+
+func (*tbMock) Helper() {}
+
+func (m *tbMock) Fatalf(string, ...any) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.fatalCalls++
+}
+
+func (m *tbMock) failed() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.fatalCalls > 0
+}
+
+func TestWaitBuffer_Bytes_ReturnsCopy(t *testing.T) {
+	t.Parallel()
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("original"))
+	require.NoError(t, err)
+
+	b := wb.Bytes()
+	// Mutate the returned slice.
+	for i := range b {
+		b[i] = 'X'
+	}
+	// The internal buffer must be unchanged.
+	require.Equal(t, "original", wb.String())
+}
+
+func TestWaitBuffer_PlainBuffer(t *testing.T) {
+	t.Parallel()
+
+	wb := testutil.NewWaitBuffer()
+	_, err := wb.Write([]byte("hello "))
+	require.NoError(t, err)
+	_, err = wb.Write([]byte("world"))
+	require.NoError(t, err)
+
+	require.Equal(t, "hello world", wb.String())
+	require.Equal(t, []byte("hello world"), wb.Bytes())
+}