74 lines
2.0 KiB
Go
74 lines
2.0 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestWorkerRuntimeStopsLoopsAndWaitsForTasks(t *testing.T) {
|
|
runtime := &workerRuntime{}
|
|
ctx, started := runtime.start()
|
|
if !started {
|
|
t.Fatal("worker runtime should start")
|
|
}
|
|
loopStopped := make(chan struct{})
|
|
if !runtime.goLoop(func(context.Context) {
|
|
<-ctx.Done()
|
|
close(loopStopped)
|
|
}) {
|
|
t.Fatal("worker loop should start")
|
|
}
|
|
taskRelease := make(chan struct{})
|
|
if !runtime.goTask(func() { <-taskRelease }) {
|
|
t.Fatal("worker task should start")
|
|
}
|
|
|
|
stopCtx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
|
defer cancel()
|
|
if err := runtime.stop(stopCtx); !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("stop should wait for the active task, got %v", err)
|
|
}
|
|
select {
|
|
case <-loopStopped:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("worker loop did not observe drain cancellation")
|
|
}
|
|
if runtime.activeTaskCount() != 1 {
|
|
t.Fatalf("active task count = %d, want 1", runtime.activeTaskCount())
|
|
}
|
|
|
|
close(taskRelease)
|
|
finalCtx, finalCancel := context.WithTimeout(context.Background(), time.Second)
|
|
defer finalCancel()
|
|
if err := runtime.stop(finalCtx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if runtime.activeTaskCount() != 0 {
|
|
t.Fatalf("active task count = %d, want 0", runtime.activeTaskCount())
|
|
}
|
|
if runtime.goTask(func() {}) {
|
|
t.Fatal("drained runtime should reject new tasks")
|
|
}
|
|
}
|
|
|
|
func TestTaskWritesAreRejectedDuringDrain(t *testing.T) {
|
|
runtime := &workerRuntime{}
|
|
runtime.start()
|
|
runtime.beginDrain()
|
|
svc := &Service{workers: runtime}
|
|
|
|
if _, err := svc.CreateTask("user-1", CreateTaskRequest{}); !isAppErrorStatus(err, 503) {
|
|
t.Fatalf("CreateTask during drain error = %v, want 503", err)
|
|
}
|
|
if _, err := svc.RetryTask("user-1", "task-1"); !isAppErrorStatus(err, 503) {
|
|
t.Fatalf("RetryTask during drain error = %v, want 503", err)
|
|
}
|
|
}
|
|
|
|
func isAppErrorStatus(err error, status int) bool {
|
|
var appErr *AppError
|
|
return errors.As(err, &appErr) && appErr.Status == status
|
|
}
|