Fix tests

internal/context/context_test.go

@@ -6,6 +6,7 @@ package context
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -61,8 +62,21 @@ func TestCreateContextWithBaseContext(t *testing.T) {
 
 	tests := []struct {
 		baseCtx  context.Context
+		validate func(ctx context.Context, t *testing.T) error
 	}{
-		{context.WithValue(context.Background(), key, value)},
+		{
+			baseCtx: context.WithValue(context.Background(), key, value),
+			validate: func(ctx context.Context, t *testing.T) error {
+				got, ok := ctx.Value(key).(ctxValue)
+				if !ok {
+					return fmt.Errorf("ctx.Value().(ctxValue) returned false, expected to be true")
+				}
+				if want := value; got != want {
+					return fmt.Errorf("ctx.Value().(ctxValue) returned unknown value (%v), expected to be %s", got, value)
+				}
+				return nil
+			},
+		},
 	}
 
 	for _, tc := range tests {
@@ -76,18 +90,13 @@ func TestCreateContextWithBaseContext(t *testing.T) {
 		defer cancel()
 
 		select {
-		case <-ctx.Done():
+		case x := <-ctx.Done():
+			t.Errorf("<-ctx.Done() == %v, want nothing (it should block)", x)
 		default:
-			t.Errorf("ctx.Done() blocked, want it to be non-blocking")
 		}
 
-		v, ok := ctx.Value(key).(ctxValue)
+		if err := tc.validate(ctx, t); err != nil {
-		original, _ := tc.baseCtx.Value(key).(ctxValue)
+			t.Errorf("%v", err)
-		if !ok {
-			t.Errorf("ctx.Value().(ctxValue) returned false, expected to be true")
-		}
-		if v != original {
-			t.Errorf("ctx.Value().(ctxValue) returned unknown value (%v), expected to be %s", v, value)
 		}
 	}
 }

processor_test.go

@@ -61,6 +61,7 @@ func newProcessorForTest(t *testing.T, r *rdb.RDB, h Handler) *processor {
 	p := newProcessor(processorParams{
 		logger:          testLogger,
 		broker:          r,
+		baseCtxFn:       context.Background,
 		retryDelayFunc:  DefaultRetryDelayFunc,
 		isFailureFunc:   defaultIsFailureFunc,
 		syncCh:          syncCh,
@@ -592,6 +593,7 @@ func TestProcessorWithStrictPriority(t *testing.T) {
 		p := newProcessor(processorParams{
 			logger:          testLogger,
 			broker:          rdbClient,
+			baseCtxFn:       context.Background,
 			retryDelayFunc:  DefaultRetryDelayFunc,
 			isFailureFunc:   defaultIsFailureFunc,
 			syncCh:          syncCh,