repo_description: Go RPC framework with high-performance and strong-extensibility for building micro-services.
repo_topics: framework go grpc microservices protobuf rpc streaming thrift
code: /*
   * Copyright 2023 CloudWeGo Authors
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package fallback
  
  import (
  	"context"
  	"errors"
  	"testing"
  
  	mockthrift "github.com/cloudwego/kitex/internal/mocks/thrift"
  	"github.com/cloudwego/kitex/internal/test"
  	"github.com/cloudwego/kitex/pkg/kerrors"
  	"github.com/cloudwego/kitex/pkg/rpcinfo"
  	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
  	"github.com/cloudwego/kitex/pkg/utils"
  )
  
  func TestNewFallbackPolicy(t *testing.T) {
  	// case0: policy is nil
  	var fbP *Policy
  	result := mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback := fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, errMock)
  	test.Assert(t, fbErr == errMock)
  	test.Assert(t, !reportAsFallback)
  
  	// case1: original error is non-nil and set result in fallback
  	fbP = NewFallbackPolicy(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		_, ok := args.(*mockthrift.MockTestArgs)
  		test.Assert(t, ok)
  		result.SetSuccess(&retStr)
  		return nil
  	})
  	test.Assert(t, !fbP.reportAsFallback)
  	test.Assert(t, fbP.fallbackFunc != nil)
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, errMock)
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil, fbErr)
  	test.Assert(t, !reportAsFallback)
  
  	// case2: enable reportAsFallback
  	fbP = NewFallbackPolicy(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		_, ok := args.(*mockthrift.MockTestArgs)
  		test.Assert(t, ok)
  		result.SetSuccess(&retStr)
  		return nil
  	}).EnableReportAsFallback()
  	test.Assert(t, fbP.reportAsFallback)
  	_, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), mockthrift.NewMockTestResult(), errMock)
  	test.Assert(t, reportAsFallback)
  
  	// case3: original error is nil, still can update result
  	fbP = NewFallbackPolicy(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		_, ok := args.(*mockthrift.MockTestArgs)
  		test.Assert(t, ok)
  		result.SetSuccess(&retStr)
  		return nil
  	}).EnableReportAsFallback()
  	test.Assert(t, fbP.reportAsFallback)
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, nil)
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, !reportAsFallback)
  
  	// case4: fallback return nil-resp and nil-err, framework will return the original resp and err
  	fbP = NewFallbackPolicy(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		_, ok := args.(*mockthrift.MockTestArgs)
  		test.Assert(t, ok)
  		return
  	}).EnableReportAsFallback()
  	test.Assert(t, fbP.reportAsFallback)
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), mockthrift.NewMockTestResult(), errMock)
  	test.Assert(t, fbErr == errMock)
  	test.Assert(t, !reportAsFallback)
  
  	// case5: WithRealReqResp, original error is non-nil and set result in fallback
  	fbP = NewFallbackPolicy(UnwrapHelper(func(ctx context.Context, req, resp interface{}, err error) (fbResp interface{}, fbErr error) {
  		_, ok := req.(*mockthrift.MockReq)
  		test.Assert(t, ok, req)
  		return &retStr, nil
  	}))
  	test.Assert(t, !fbP.reportAsFallback)
  	test.Assert(t, fbP.fallbackFunc != nil)
  	args := mockthrift.NewMockTestArgs()
  	args.Req = &mockthrift.MockReq{}
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), args, result, errMock)
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, !reportAsFallback)
  }
  
  func TestErrorFallback(t *testing.T) {
  	// case1: err is non-nil
  	fbP := ErrorFallback(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		result.SetSuccess(&retStr)
  		return nil
  	})
  	result := mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback := fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, errMock)
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, !reportAsFallback)
  
  	// case 2: err is nil, then the fallback func won't be executed
  	fbExecuted := false
  	fbP = ErrorFallback(UnwrapHelper(func(ctx context.Context, req, resp interface{}, err error) (fbResp interface{}, fbErr error) {
  		fbExecuted = true
  		return &retStr, nil
  	})).EnableReportAsFallback()
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, nil)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, !reportAsFallback)
  	test.Assert(t, !fbExecuted)
  }
  
  func TestTimeoutAndCBFallback(t *testing.T) {
  	// case1: rpc timeout will do fallback
  	fbP := TimeoutAndCBFallback(UnwrapHelper(func(ctx context.Context, req, resp interface{}, err error) (fbResp interface{}, fbErr error) {
  		_, ok := req.(*mockthrift.MockReq)
  		test.Assert(t, ok)
  		return &retStr, nil
  	}))
  	result := mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback := fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, kerrors.ErrRPCTimeout.WithCause(errMock))
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, !reportAsFallback)
  
  	// case2: circuit breaker error will do fallback
  	fbP = TimeoutAndCBFallback(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		result.SetSuccess(&retStr)
  		return nil
  	}).EnableReportAsFallback()
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, kerrors.ErrCircuitBreak.WithCause(errMock))
  	test.Assert(t, *result.GetResult().(*string) == retStr)
  	test.Assert(t, fbErr == nil)
  	test.Assert(t, reportAsFallback)
  
  	// case3: err is non-nil, but not rpc timeout or circuit breaker
  	fbP = TimeoutAndCBFallback(func(ctx context.Context, args utils.KitexArgs, result utils.KitexResult, err error) (fbErr error) {
  		result.SetSuccess(&retStr)
  		return nil
  	})
  	result = mockthrift.NewMockTestResult()
  	fbErr, reportAsFallback = fbP.DoIfNeeded(context.Background(), genRPCInfo(), mockthrift.NewMockTestArgs(), result, errMock)
  	test.Assert(t, fbErr == errMock)
  	test.Assert(t, !reportAsFallback)
  }
  
  var (
  	method  = "test"
  	errMock = errors.New("mock")
  	retStr  = "success"
  )
  
  func genRPCInfo() rpcinfo.RPCInfo {
  	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{Method: method}, method).ImmutableView()
  	ri := rpcinfo.NewRPCInfo(to, to, rpcinfo.NewInvocation("", method), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
  	return ri
  }
  