test(queue): cover handler context propagation across drivers

Author: cmilesio

driver/redisqueue/worker_redis_impl.go

@@ -20,6 +20,7 @@ type redisWorker struct {
 	server server
 	mux    *backend.ServeMux
 	obs    queue.Observer
+	ctxDecorator func(context.Context) context.Context
 
 	mu      sync.Mutex
 	started bool
@@ -29,17 +30,31 @@ func newRedisWorker(server server, mux *backend.ServeMux, observer queue.Observe
 	return &redisWorker{server: server, mux: mux, obs: observer}
 }
 
+func (w *redisWorker) SetHandlerContextDecorator(fn func(context.Context) context.Context) {
+	w.ctxDecorator = fn
+}
+
 func (w *redisWorker) Register(jobType string, handler queue.Handler) {
 	if jobType == "" || handler == nil {
 		return
 	}
 	if w.obs == nil {
 		w.mux.HandleFunc(jobType, func(ctx context.Context, job *backend.Task) error {
+			if w.ctxDecorator != nil {
+				if decorated := w.ctxDecorator(ctx); decorated != nil {
+					ctx = decorated
+				}
+			}
 			return handler(ctx, queue.NewJob(job.Type()).Payload(job.Payload()))
 		})
 		return
 	}
 	w.mux.HandleFunc(jobType, func(ctx context.Context, job *backend.Task) error {
+		if w.ctxDecorator != nil {
+			if decorated := w.ctxDecorator(ctx); decorated != nil {
+				ctx = decorated
+			}
+		}
 		attempt, _ := backend.GetRetryCount(ctx)
 		maxRetry, _ := backend.GetMaxRetry(ctx)
 		queueName, _ := backend.GetQueueName(ctx)

driver/redisqueue/worker_redis_impl_test.go

@@ -101,7 +101,7 @@ func TestRedisWorker_ShutdownHonorsContext(t *testing.T) {
 func TestRedisWorker_ProcessEventsWithObserver(t *testing.T) {
 	server := &serverStub{}
 	var events []queue.Event
-	observer := queue.ObserverFunc(func(event queue.Event) { events = append(events, event) })
+	observer := queue.ObserverFunc(func(_ context.Context, event queue.Event) { events = append(events, event) })
 	w := newRedisWorker(server, backend.NewServeMux(), observer)
 
 	w.Register("job:ok", func(context.Context, queue.Job) error { return nil })
@@ -156,7 +156,7 @@ func TestRedisWorker_ProcessEventsWithObserver(t *testing.T) {
 func TestRedisWorker_ProcessEventsUnwrapBusEnvelopeJobType(t *testing.T) {
 	server := &serverStub{}
 	var events []queue.Event
-	observer := queue.ObserverFunc(func(event queue.Event) { events = append(events, event) })
+	observer := queue.ObserverFunc(func(_ context.Context, event queue.Event) { events = append(events, event) })
 	w := newRedisWorker(server, backend.NewServeMux(), observer)
 
 	w.Register("bus:job", func(context.Context, queue.Job) error { return nil })
@@ -210,3 +210,48 @@ func TestRedisWorker_NoObserverFastPath(t *testing.T) {
 		t.Fatalf("expected handler called once, got %d", called)
 	}
 }
+
+func TestRedisWorker_ObserverSeesDecoratedContext(t *testing.T) {
+	server := &serverStub{}
+	type ctxKey struct{}
+	key := ctxKey{}
+	const want = "jobs"
+
+	var observed []string
+	var handled []string
+	observer := queue.ObserverFunc(func(ctx context.Context, event queue.Event) {
+		if event.Kind != queue.EventProcessStarted && event.Kind != queue.EventProcessSucceeded {
+			return
+		}
+		value, _ := ctx.Value(key).(string)
+		observed = append(observed, value)
+	})
+	w := newRedisWorker(server, backend.NewServeMux(), observer)
+	w.SetHandlerContextDecorator(func(ctx context.Context) context.Context {
+		return context.WithValue(ctx, key, want)
+	})
+
+	w.Register("job:decorated", func(ctx context.Context, _ queue.Job) error {
+		value, _ := ctx.Value(key).(string)
+		handled = append(handled, value)
+		return nil
+	})
+	if err := w.StartWorkers(context.Background()); err != nil {
+		t.Fatalf("start workers failed: %v", err)
+	}
+	if err := server.lastStartHandler.ProcessTask(context.Background(), backend.NewTask("job:decorated", []byte("ok"))); err != nil {
+		t.Fatalf("process task failed: %v", err)
+	}
+
+	if len(observed) != 2 {
+		t.Fatalf("expected 2 observed events, got %d", len(observed))
+	}
+	for i, got := range observed {
+		if got != want {
+			t.Fatalf("expected observed[%d] = %q, got %q", i, want, got)
+		}
+	}
+	if len(handled) != 1 || handled[0] != want {
+		t.Fatalf("expected handler to see %q, got %#v", want, handled)
+	}
+}