fix unit test

Signed-off-by: Shijie Sheng <liouvetren@gmail.com>

Author: shijiesheng

common/rpc/middleware.go

@@ -23,12 +23,10 @@ package rpc
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"io"
 
 	"go.uber.org/cadence/worker"
 	"go.uber.org/yarpc"
-	"go.uber.org/yarpc/api/middleware"
 	"go.uber.org/yarpc/api/transport"
 
 	"github.com/uber/cadence/common"

service/frontend/wrappers/clusterredirection/utils_test.go

@@ -25,47 +25,51 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.uber.org/yarpc/yarpctest"
 
+	apiv1 "github.com/uber/cadence-idl/go/proto/api/v1"
 	"github.com/uber/cadence/common"
+	"github.com/uber/cadence/common/client"
 	"github.com/uber/cadence/common/types"
 )
 
 func TestGetRequestedConsistencyLevelFromContext(t *testing.T) {
 	tests := []struct {
 		name     string
-		ctx      context.Context
+		featureFlags apiv1.FeatureFlags
 		expected types.QueryConsistencyLevel
 	}{
 		{
-			name:     "no value in context returns eventual",
-			ctx:      context.Background(),
+			name:     "empty feature flags",
+			featureFlags: apiv1.FeatureFlags{},
 			expected: types.QueryConsistencyLevelEventual,
 		},
 		{
-			name:     "eventual level in context",
-			ctx:      context.WithValue(context.Background(), common.QueryConsistencyLevelHeaderName, types.QueryConsistencyLevelEventual),
+			name:     "auto forwarding disabled",
+			featureFlags: apiv1.FeatureFlags{AutoforwardingEnabled: false},
 			expected: types.QueryConsistencyLevelEventual,
 		},
 		{
-			name:     "strong level in context",
-			ctx:      context.WithValue(context.Background(), common.QueryConsistencyLevelHeaderName, types.QueryConsistencyLevelStrong),
+			name:     "autoforwarding enabled",
+			featureFlags: apiv1.FeatureFlags{AutoforwardingEnabled: true},
 			expected: types.QueryConsistencyLevelStrong,
 		},
 		{
-			name:     "wrong type in context returns eventual",
-			ctx:      context.WithValue(context.Background(), common.QueryConsistencyLevelHeaderName, "strong"),
-			expected: types.QueryConsistencyLevelEventual,
-		},
-		{
-			name:     "nil context value returns eventual",
-			ctx:      context.WithValue(context.Background(), common.QueryConsistencyLevelHeaderName, nil),
+			name:     "no autoforwarding field",
+			featureFlags: apiv1.FeatureFlags{
+				WorkflowExecutionAlreadyCompletedErrorEnabled: true,
+			},
 			expected: types.QueryConsistencyLevelEventual,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			result := getRequestedConsistencyLevelFromContext(tt.ctx)
+			ctx := yarpctest.ContextWithCall(context.Background(), &yarpctest.Call{
+				Headers: map[string]string{common.ClientFeatureFlagsHeaderName: client.FeatureFlagsHeader(tt.featureFlags)},
+			})
+
+			result := getRequestedConsistencyLevelFromContext(ctx)
 			assert.Equal(t, tt.expected, result)
 		})
 	}