253 lines
6.7 KiB
Go
253 lines
6.7 KiB
Go
package util
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
// Keep worker log lines out of the test output.
|
|
slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, nil)))
|
|
os.Exit(m.Run())
|
|
}
|
|
|
|
func TestWorkerRunRejectsBadConfig(t *testing.T) {
|
|
noop := func(context.Context, time.Time) error { return nil }
|
|
tests := []struct {
|
|
name string
|
|
w Worker
|
|
}{
|
|
{"no task", Worker{Name: "a", Interval: time.Second}},
|
|
{"zero interval", Worker{Name: "b", Task: noop}},
|
|
{"negative interval", Worker{Name: "c", Interval: -time.Second, Task: noop}},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
if err := tc.w.Run(context.Background(), time.Time{}); err == nil {
|
|
t.Fatal("expected an error")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWorkerRunCancelledBeforeStart(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
called := false
|
|
w := Worker{Name: "w", Interval: time.Millisecond, Task: func(context.Context, time.Time) error {
|
|
called = true
|
|
return nil
|
|
}}
|
|
if err := w.Run(ctx, time.Time{}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if called {
|
|
t.Fatal("task ran after the context was already cancelled")
|
|
}
|
|
}
|
|
|
|
// The checkpoint passed to each run advances only when the previous run
|
|
// succeeded or reported ErrWorkerNoWork.
|
|
func TestWorkerRunCheckpoints(t *testing.T) {
|
|
errFail := errors.New("boom")
|
|
results := []error{
|
|
nil, // advance
|
|
errFail, // keep
|
|
context.Canceled, // keep
|
|
fmt.Errorf("%w: x", ErrWorkerNoWork), // advance
|
|
nil,
|
|
}
|
|
since0 := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
var (
|
|
mu sync.Mutex
|
|
seen []time.Time
|
|
starts []time.Time
|
|
)
|
|
w := Worker{
|
|
Name: "w",
|
|
Interval: 5 * time.Millisecond,
|
|
Task: func(_ context.Context, since time.Time) error {
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
i := len(seen)
|
|
seen = append(seen, since)
|
|
starts = append(starts, time.Now())
|
|
if i == len(results)-1 {
|
|
cancel()
|
|
}
|
|
return results[i]
|
|
},
|
|
}
|
|
|
|
done := make(chan error, 1)
|
|
go func() { done <- w.Run(ctx, since0) }()
|
|
select {
|
|
case err := <-done:
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("worker did not stop")
|
|
}
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(seen) != len(results) {
|
|
t.Fatalf("task ran %d times, want %d", len(seen), len(results))
|
|
}
|
|
if !seen[0].Equal(since0) {
|
|
t.Fatalf("first run since = %v, want %v", seen[0], since0)
|
|
}
|
|
// Run 0 succeeded, so runs 1..3 see a checkpoint from run 0's start.
|
|
for i := 1; i <= 3; i++ {
|
|
if !seen[i].Equal(seen[1]) {
|
|
t.Fatalf("run %d since = %v, want unchanged %v", i, seen[i], seen[1])
|
|
}
|
|
}
|
|
if !seen[1].After(since0) || seen[1].After(starts[1]) {
|
|
t.Fatalf("checkpoint after success = %v, want between %v and %v", seen[1], since0, starts[1])
|
|
}
|
|
// Run 3 returned ErrWorkerNoWork, so run 4 sees a newer checkpoint.
|
|
if !seen[4].After(seen[3]) {
|
|
t.Fatalf("checkpoint did not advance after ErrWorkerNoWork: %v -> %v", seen[3], seen[4])
|
|
}
|
|
}
|
|
|
|
// An in-flight task keeps a live context after the parent is cancelled and is
|
|
// allowed to finish.
|
|
func TestWorkerRunGracefulShutdown(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
var workErr error
|
|
w := Worker{
|
|
Name: "w", Interval: time.Hour, ShutdownTimeout: time.Second,
|
|
Task: func(taskCtx context.Context, _ time.Time) error {
|
|
cancel()
|
|
time.Sleep(10 * time.Millisecond)
|
|
workErr = taskCtx.Err()
|
|
return nil
|
|
},
|
|
}
|
|
if err := w.Run(ctx, time.Time{}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if workErr != nil {
|
|
t.Fatalf("task context cancelled before shutdown timeout: %v", workErr)
|
|
}
|
|
}
|
|
|
|
// A task that outlives ShutdownTimeout has its context cancelled.
|
|
func TestWorkerRunShutdownTimeout(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
w := Worker{
|
|
Name: "w", Interval: time.Hour, ShutdownTimeout: 20 * time.Millisecond,
|
|
Task: func(taskCtx context.Context, _ time.Time) error {
|
|
cancel()
|
|
select {
|
|
case <-taskCtx.Done():
|
|
return taskCtx.Err()
|
|
case <-time.After(5 * time.Second):
|
|
return errors.New("task context never cancelled")
|
|
}
|
|
},
|
|
}
|
|
start := time.Now()
|
|
if err := w.Run(ctx, time.Time{}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if d := time.Since(start); d > 2*time.Second {
|
|
t.Fatalf("shutdown took %v", d)
|
|
}
|
|
}
|
|
|
|
func TestWorkerExecute(t *testing.T) {
|
|
since := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
|
logger := slog.Default()
|
|
tests := []struct {
|
|
name string
|
|
err error
|
|
advance bool
|
|
}{
|
|
{"success", nil, true},
|
|
{"no work", ErrWorkerNoWork, true},
|
|
{"wrapped no work", fmt.Errorf("%w: nothing", ErrWorkerNoWork), true},
|
|
{"cancelled", context.Canceled, false},
|
|
{"wrapped cancelled", fmt.Errorf("x: %w", context.Canceled), false},
|
|
{"failure", errors.New("boom"), false},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
w := Worker{Task: func(context.Context, time.Time) error { return tc.err }}
|
|
before := time.Now()
|
|
got := w.execute(context.Background(), logger, since)
|
|
if tc.advance {
|
|
if got.Before(before) {
|
|
t.Fatalf("checkpoint %v not advanced to run start (>= %v)", got, before)
|
|
}
|
|
} else if !got.Equal(since) {
|
|
t.Fatalf("checkpoint = %v, want unchanged %v", got, since)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAwaitTask(t *testing.T) {
|
|
errFail := errors.New("boom")
|
|
v, err := AwaitTask(context.Background(), func() (int, error) { return 42, nil })
|
|
if v != 42 || err != nil {
|
|
t.Fatalf("got (%d, %v), want (42, nil)", v, err)
|
|
}
|
|
v, err = AwaitTask(context.Background(), func() (int, error) { return 7, errFail })
|
|
if v != 7 || !errors.Is(err, errFail) {
|
|
t.Fatalf("got (%d, %v), want (7, %v)", v, err, errFail)
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
release := make(chan struct{})
|
|
defer close(release)
|
|
go func() { time.Sleep(5 * time.Millisecond); cancel() }()
|
|
v, err = AwaitTask(ctx, func() (int, error) { <-release; return 1, nil })
|
|
if v != 0 || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("got (%d, %v), want (0, context.Canceled)", v, err)
|
|
}
|
|
}
|
|
|
|
func TestAwaitTaskE(t *testing.T) {
|
|
errFail := errors.New("boom")
|
|
if err := AwaitTaskE(context.Background(), func() error { return nil }); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := AwaitTaskE(
|
|
context.Background(),
|
|
func() error { return errFail },
|
|
); !errors.Is(
|
|
err,
|
|
errFail,
|
|
) {
|
|
t.Fatalf("got %v, want %v", err, errFail)
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond)
|
|
defer cancel()
|
|
release := make(chan struct{})
|
|
defer close(release)
|
|
if err := AwaitTaskE(
|
|
ctx,
|
|
func() error { <-release; return nil },
|
|
); !errors.Is(
|
|
err,
|
|
context.DeadlineExceeded,
|
|
) {
|
|
t.Fatalf("got %v, want context.DeadlineExceeded", err)
|
|
}
|
|
}
|