From 5490d2c6257a86600ea5b198b97b27e6e0ab076f Mon Sep 17 00:00:00 2001
From: Ken Hibino <ken.hibino7@gmail.com>
Date: Wed, 16 Feb 2022 06:51:09 -0800
Subject: [PATCH] Fix tests

---
 internal/context/context_test.go | 31 ++++++++++++++++++++-----------
 processor_test.go                |  2 ++
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/internal/context/context_test.go b/internal/context/context_test.go
index 3befdba..5e67cc8 100644
--- a/internal/context/context_test.go
+++ b/internal/context/context_test.go
@@ -6,6 +6,7 @@ package context
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -60,9 +61,22 @@ func TestCreateContextWithBaseContext(t *testing.T) {
 	var value ctxValue = "value"
 
 	tests := []struct {
-		baseCtx context.Context
+		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)
-		original, _ := tc.baseCtx.Value(key).(ctxValue)
-		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)
+		if err := tc.validate(ctx, t); err != nil {
+			t.Errorf("%v", err)
 		}
 	}
 }
diff --git a/processor_test.go b/processor_test.go
index 21ba01f..5d3db46 100644
--- a/processor_test.go
+++ b/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,