/*
 * Copyright 2021 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 retry

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/stats"
)

func BenchmarkRandomBackOff_Wait(b *testing.B) {
	min := 10
	max := 100
	bk := &randomBackOff{
		minMS: min,
		maxMS: max,
	}
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			ms := int(bk.randomDuration().Milliseconds())
			test.Assert(b, ms <= max && ms >= min, ms)
		}
	})
}

func TestRandomBackOff_Wait(t *testing.T) {
	const (
		min = 99
		max = 100
	)
	bk := &randomBackOff{
		minMS: min,
		maxMS: max,
	}
	startTime := time.Now()
	bk.Wait(1)
	waitTime := time.Since(startTime)
	test.Assert(t, time.Millisecond*min <= waitTime)
	test.Assert(t, waitTime < time.Millisecond*(max+30))
}

func TestRandomBackOff_String(t *testing.T) {
	min := 10
	max := 100
	bk := &randomBackOff{
		minMS: min,
		maxMS: max,
	}
	msg := "RandomBackOff(10ms-100ms)"
	test.Assert(t, bk.String() == msg)
}

func TestFixedBackOff_Wait(t *testing.T) {
	const fix = 50
	bk := &fixedBackOff{
		fixMS: fix,
	}
	startTime := time.Now()
	bk.Wait(1)
	waitTime := time.Since(startTime)
	test.Assert(t, time.Millisecond*fix <= waitTime)
	test.Assert(t, waitTime < time.Millisecond*(fix*2))
}

func TestFixedBackOff_String(t *testing.T) {
	fix := 50
	bk := &fixedBackOff{
		fixMS: fix,
	}
	msg := "FixedBackOff(50ms)"
	test.Assert(t, bk.String() == msg)
}

func TestNoneBackOff_String(t *testing.T) {
	bk := &noneBackOff{}
	msg := "NoneBackOff"
	test.Assert(t, bk.String() == msg)
}

// test FailurePolicy call
func TestFailurePolicyCall(t *testing.T) {
	// call while rpc timeout
	ctx := context.Background()
	rc := NewRetryContainer()
	failurePolicy := NewFailurePolicy()
	failurePolicy.BackOffPolicy.BackOffType = FixedBackOffType
	failurePolicy.BackOffPolicy.CfgItems = map[BackOffCfgKey]float64{
		FixMSBackOffCfgKey: 100.0,
	}
	failurePolicy.StopPolicy.MaxDurationMS = 100
	err := rc.Init(map[string]Policy{Wildcard: {
		Enable:        true,
		Type:          0,
		FailurePolicy: failurePolicy,
	}}, nil)
	test.Assert(t, err == nil, err)
	ri := genRPCInfo()
	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
	_, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		return ri, kerrors.ErrRPCTimeout
	}, ri, nil, nil)
	test.Assert(t, err != nil, err)
	test.Assert(t, !ok)

	// call normal
	failurePolicy.StopPolicy.MaxDurationMS = 0
	err = rc.Init(map[string]Policy{Wildcard: {
		Enable:        true,
		Type:          0,
		FailurePolicy: failurePolicy,
	}}, nil)
	test.Assert(t, err == nil, err)
	_, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		return ri, nil
	}, ri, nil, nil)
	test.Assert(t, err == nil, err)
	test.Assert(t, ok)
}

// test retry with one time policy
func TestRetryWithOneTimePolicy(t *testing.T) {
	// call while rpc timeout and exceed MaxDurationMS cause BackOffPolicy is wait fix 100ms, it is invalid config
	failurePolicy := NewFailurePolicy()
	failurePolicy.BackOffPolicy.BackOffType = FixedBackOffType
	failurePolicy.BackOffPolicy.CfgItems = map[BackOffCfgKey]float64{
		FixMSBackOffCfgKey: 100.0,
	}
	failurePolicy.StopPolicy.MaxDurationMS = 100
	p := Policy{
		Enable:        true,
		Type:          0,
		FailurePolicy: failurePolicy,
	}
	ri := genRPCInfo()
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	_, ok, err := NewRetryContainer().WithRetryIfNeeded(ctx, &p, func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		return ri, kerrors.ErrRPCTimeout
	}, ri, nil, nil)
	test.Assert(t, err != nil, err)
	test.Assert(t, !ok)

	// call no MaxDurationMS limit, the retry will success
	failurePolicy.StopPolicy.MaxDurationMS = 0
	var callTimes int32
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), genRPCInfo())
	_, ok, err = NewRetryContainer().WithRetryIfNeeded(ctx, &p, func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		if atomic.LoadInt32(&callTimes) == 0 {
			atomic.AddInt32(&callTimes, 1)
			return ri, kerrors.ErrRPCTimeout
		}
		return ri, nil
	}, ri, nil, nil)
	test.Assert(t, err == nil, err)
	test.Assert(t, !ok)

	// call backup request
	p = BuildBackupRequest(NewBackupPolicy(10))
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), genRPCInfo())
	callTimes = 0
	_, ok, err = NewRetryContainer().WithRetryIfNeeded(ctx, &p, func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		if atomic.LoadInt32(&callTimes) == 0 || atomic.LoadInt32(&callTimes) == 1 {
			atomic.AddInt32(&callTimes, 1)
			time.Sleep(time.Millisecond * 100)
		}
		return ri, nil
	}, ri, nil, nil)
	test.Assert(t, err == nil, err)
	test.Assert(t, !ok)
	test.Assert(t, atomic.LoadInt32(&callTimes) == 2)
}

// test specified error to retry
func TestSpecifiedErrorRetry(t *testing.T) {
	ri := genRPCInfo()
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	var transErrCode int32 = 1000

	// case1: specified method retry with error
	t.Run("case1", func(t *testing.T) {
		rc := NewRetryContainer()
		checkResultRetry := false
		shouldResultRetry := &ShouldResultRetry{ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() == method {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == transErrCode {
					checkResultRetry = true
					return true
				}
			}
			return false
		}}
		err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, checkResultRetry)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})

	// case2: specified method retry with error, but use backup request config cannot be effective
	t.Run("case2", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() == method {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == transErrCode {
					return true
				}
			}
			return false
		}}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{Wildcard: BuildBackupRequest(NewBackupPolicy(10))}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri = genRPCInfo()
		_, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err != nil, err)
		test.Assert(t, !ok)
	})

	// case3: specified method retry with error, but method not match
	t.Run("case3", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() != method {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == transErrCode {
					return true
				}
			}
			return false
		}}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{method: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri = genRPCInfo()
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err != nil)
		test.Assert(t, !ok)
		_, ok = ri.To().Tag(remoteTagKey)
		test.Assert(t, !ok)
	})

	// case4: all error retry
	t.Run("case4", func(t *testing.T) {
		rc := NewRetryContainer()
		p := BuildFailurePolicy(NewFailurePolicyWithResultRetry(AllErrorRetry()))
		ri = genRPCInfo()
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &p, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})

	// case5: all error retry and trigger circuit breaker
	t.Run("case5", func(t *testing.T) {
		rc := NewRetryContainerWithPercentageLimit()
		p := BuildFailurePolicy(NewFailurePolicyWithResultRetry(AllErrorRetry()))
		ri = genRPCInfo()

		for i := 0; i < 10; i++ {
			// failure rate is 50%
			_, _, err := rc.WithRetryIfNeeded(ctx, &p, retryWithTransError(0, transErrCode), ri, nil, nil)
			if i < 5 {
				// i < 5, total request < 10
				test.Assert(t, err == nil, err, i)
			} else {
				// the minSimple is 10, total request > 10 then trigger cb
				test.Assert(t, err != nil)
			}
		}
	})
}

// test specified resp to retry
func TestSpecifiedRespRetry(t *testing.T) {
	retryResult := &mockResult{}
	retryResp := mockResp{
		code: 500,
		msg:  "retry",
	}
	noRetryResp := mockResp{
		code: 0,
		msg:  "noretry",
	}
	var callTimes int32
	retryWithResp := func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		retryResult := response.(*mockResult)
		newVal := atomic.AddInt32(&callTimes, 1)
		if newVal == 1 {
			retryResult.setResult(retryResp)
			return genRPCInfo(), nil
		} else {
			retryResult.setResult(noRetryResp)
			return genRPCInfoWithRemoteTag(remoteTags), nil
		}
	}
	ctx := context.Background()
	ri := genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
	rc := NewRetryContainer()
	// case1: specified method retry with resp
	shouldResultRetry := &ShouldResultRetry{RespRetry: func(resp interface{}, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() == method {
			if r, ok := resp.(*mockResult); ok && r.getResult() == retryResp {
				return true
			}
		}
		return false
	}}
	err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == noRetryResp, retryResult)
	test.Assert(t, !ok)
	v, ok := ri.To().Tag(remoteTagKey)
	test.Assert(t, ok)
	test.Assert(t, v == remoteTagValue)

	// case2 specified method retry with resp, but use backup request config cannot be effective
	retryResult = &mockResult{}
	atomic.StoreInt32(&callTimes, 0)
	rc = NewRetryContainer()
	err = rc.Init(map[string]Policy{Wildcard: BuildBackupRequest(NewBackupPolicy(100))}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri = genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	_, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == retryResp, retryResp)
	test.Assert(t, !ok)

	// case3: specified method retry with resp, but method not match
	atomic.StoreInt32(&callTimes, 0)
	shouldResultRetry = &ShouldResultRetry{RespRetry: func(resp interface{}, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() != method {
			if r, ok := resp.(*mockResult); ok && r.getResult() == retryResp {
				return true
			}
		}
		return false
	}}
	rc = NewRetryContainer()
	err = rc.Init(map[string]Policy{method: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri = genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	ri, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == retryResp, retryResult)
	test.Assert(t, ok)
	_, ok = ri.To().Tag(remoteTagKey)
	test.Assert(t, !ok)
}

// test specified error to retry with ErrorRetryWithCtx
func TestSpecifiedErrorRetryWithCtx(t *testing.T) {
	ri := genRPCInfo()
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	var transErrCode int32 = 1000

	// case1: specified method retry with error
	t.Run("case1", func(t *testing.T) {
		rc := NewRetryContainer()
		shouldResultRetry := &ShouldResultRetry{ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() == method {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == transErrCode {
					return true
				}
			}
			return false
		}}
		err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})

	// case2: specified method retry with error, but use backup request config cannot be effective
	t.Run("case2", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() == method {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == transErrCode {
					return true
				}
			}
			return false
		}}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{Wildcard: BuildBackupRequest(NewBackupPolicy(10))}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri = genRPCInfo()
		_, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err != nil, err)
		test.Assert(t, !ok)
	})

	// case3: specified method retry with error, but method not match
	t.Run("case3", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
			return ri.To().Method() != method
		}}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{method: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri = genRPCInfo()
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err != nil)
		test.Assert(t, !ok)
		_, ok = ri.To().Tag(remoteTagKey)
		test.Assert(t, !ok)
	})

	// case4: all error retry
	t.Run("case4", func(t *testing.T) {
		rc := NewRetryContainer()
		p := BuildFailurePolicy(NewFailurePolicyWithResultRetry(AllErrorRetry()))
		ri = genRPCInfo()
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &p, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})

	// case5: specified method retry with error, only ctx has some info then retry
	ctxKeyVal := "ctxKeyVal"
	t.Run("case5", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
			if ri.To().Method() == method && ctx.Value(ctxKeyVal) == ctxKeyVal {
				if te, ok := err.(*remote.TransError); ok && te.TypeID() == 1000 {
					return true
				}
			}
			return false
		}}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{method: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri = genRPCInfo()
		ctx = context.WithValue(ctx, ctxKeyVal, ctxKeyVal)
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, transErrCode), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})
}

// test specified error to retry, but has both old and new policy, the new one will be effective
func TestSpecifiedErrorRetryHasOldAndNew(t *testing.T) {
	ri := genRPCInfo()
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	// case1: ErrorRetryWithCtx will retry, but ErrorRetry not retry, the expect result is do retry
	t.Run("case1", func(t *testing.T) {
		rc := NewRetryContainer()
		shouldResultRetry := &ShouldResultRetry{
			ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
				return true
			},
			ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
				return false
			},
		}
		err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, 1000), ri, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, !ok)
		v, ok := ri.To().Tag(remoteTagKey)
		test.Assert(t, ok)
		test.Assert(t, v == remoteTagValue)
	})

	// case2: ErrorRetryWithCtx not retry, but ErrorRetry retry, the expect result is that not do retry
	t.Run("case2", func(t *testing.T) {
		shouldResultRetry := &ShouldResultRetry{
			ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
				return false
			},
			ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
				return true
			},
		}
		rc := NewRetryContainer()
		err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
		test.Assert(t, err == nil, err)
		ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithTransError(0, 1000), ri, nil, nil)
		test.Assert(t, err != nil)
		test.Assert(t, !ok)
		_, ok = ri.To().Tag(remoteTagKey)
		test.Assert(t, !ok)
	})
}

// test specified resp to retry with ErrorRetryWithCtx
func TestSpecifiedRespRetryWithCtx(t *testing.T) {
	retryResult := &mockResult{}
	retryResp := mockResp{
		code: 500,
		msg:  "retry",
	}
	noRetryResp := mockResp{
		code: 0,
		msg:  "noretry",
	}
	var callTimes int32
	retryWithResp := func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		newVal := atomic.AddInt32(&callTimes, 1)
		retryResult := response.(*mockResult)
		if newVal == 1 {
			retryResult.setResult(retryResp)
			return genRPCInfo(func() interface{} {
				return &mockResult{}
			}), nil
		} else {
			retryResult.setResult(noRetryResp)
			return genRPCInfoWithRemoteTag(remoteTags), nil
		}
	}
	ctx := context.Background()
	ri := genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
	rc := NewRetryContainer()
	// case1: specified method retry with resp
	shouldResultRetry := &ShouldResultRetry{RespRetryWithCtx: func(ctx context.Context, resp interface{}, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() == method {
			if r, ok := resp.(*mockResult); ok && r.getResult() == retryResp {
				return true
			}
		}
		return false
	}}
	err := rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri, ok, err := rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == noRetryResp, retryResult)
	test.Assert(t, !ok)
	v, ok := ri.To().Tag(remoteTagKey)
	test.Assert(t, ok)
	test.Assert(t, v == remoteTagValue)

	// case2 specified method retry with resp, but use backup request config cannot be effective
	atomic.StoreInt32(&callTimes, 0)
	rc = NewRetryContainer()
	err = rc.Init(map[string]Policy{Wildcard: BuildBackupRequest(NewBackupPolicy(100))}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri = genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	_, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == retryResp, retryResp)
	test.Assert(t, !ok)

	// case3: specified method retry with resp, but method not match
	atomic.StoreInt32(&callTimes, 0)
	shouldResultRetry = &ShouldResultRetry{RespRetryWithCtx: func(ctx context.Context, resp interface{}, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() != method {
			if r, ok := resp.(*mockResult); ok && r.getResult() == retryResp {
				return true
			}
		}
		return false
	}}
	rc = NewRetryContainer()
	err = rc.Init(map[string]Policy{method: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry)
	test.Assert(t, err == nil, err)
	ri = genRPCInfo(func() interface{} {
		return &mockResult{}
	})
	ctx = rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
	ri, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == retryResp, retryResult)
	test.Assert(t, ok)
	_, ok = ri.To().Tag(remoteTagKey)
	test.Assert(t, !ok)

	// case4: specified method retry with resp, only ctx has some info then retry
	ctxKeyVal := "ctxKeyVal"
	atomic.StoreInt32(&callTimes, 0)
	shouldResultRetry2 := &ShouldResultRetry{RespRetryWithCtx: func(ctx context.Context, resp interface{}, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() == method && ctx.Value(ctxKeyVal) == ctxKeyVal {
			if r, ok := resp.(*mockResult); ok && r.getResult() == retryResp {
				return true
			}
		}
		return false
	}}
	err = rc.Init(map[string]Policy{Wildcard: BuildFailurePolicy(NewFailurePolicy())}, shouldResultRetry2)
	test.Assert(t, err == nil, err)
	ctx = context.WithValue(ctx, ctxKeyVal, ctxKeyVal)
	ri, ok, err = rc.WithRetryIfNeeded(ctx, &Policy{}, retryWithResp, ri, nil, retryResult)
	test.Assert(t, err == nil, err)
	test.Assert(t, retryResult.getResult() == noRetryResp, retryResult)
	test.Assert(t, !ok)
	v, ok = ri.To().Tag(remoteTagKey)
	test.Assert(t, ok)
	test.Assert(t, v == remoteTagValue)
}

func TestResultRetryWithPolicyChange(t *testing.T) {
	rc := NewRetryContainer()
	shouldResultRetry := &ShouldResultRetry{ErrorRetry: func(err error, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() == method {
			if te, ok := err.(*remote.TransError); ok && te.TypeID() == 1000 {
				return true
			}
		}
		return false
	}}
	err := rc.Init(nil, shouldResultRetry)
	test.Assert(t, err == nil, err)

	// case 1: first time trigger NotifyPolicyChange, the `initRetryer` will be executed, check if the ShouldResultRetry is not nil
	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
	r := rc.getRetryer(context.Background(), genRPCInfo())
	fr, ok := r.(*failureRetryer)
	test.Assert(t, ok)
	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)

	// case 2: second time trigger NotifyPolicyChange, the `UpdatePolicy` will be executed, check if the ShouldResultRetry is not nil
	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
	r = rc.getRetryer(context.Background(), genRPCInfo())
	fr, ok = r.(*failureRetryer)
	test.Assert(t, ok)
	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)
}

func TestResultRetryWithCtxWhenPolicyChange(t *testing.T) {
	rc := NewRetryContainer()
	shouldResultRetry := &ShouldResultRetry{ErrorRetryWithCtx: func(ctx context.Context, err error, ri rpcinfo.RPCInfo) bool {
		if ri.To().Method() == method {
			if te, ok := err.(*remote.TransError); ok && te.TypeID() == 1000 {
				return true
			}
		}
		return false
	}}
	err := rc.Init(nil, shouldResultRetry)
	test.Assert(t, err == nil, err)

	// case 1: first time trigger NotifyPolicyChange, the `initRetryer` will be executed, check if the ShouldResultRetry is not nil
	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
	r := rc.getRetryer(context.Background(), genRPCInfo())
	fr, ok := r.(*failureRetryer)
	test.Assert(t, ok)
	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)

	// case 2: second time trigger NotifyPolicyChange, the `UpdatePolicy` will be executed, check if the ShouldResultRetry is not nil
	rc.NotifyPolicyChange(Wildcard, BuildFailurePolicy(NewFailurePolicy()))
	r = rc.getRetryer(context.Background(), genRPCInfo())
	fr, ok = r.(*failureRetryer)
	test.Assert(t, ok)
	test.Assert(t, fr.policy.ShouldResultRetry == shouldResultRetry)
}

func TestFailureRetryWithRPCInfo(t *testing.T) {
	// failure retry
	ctx := context.Background()
	rc := NewRetryContainer()

	ri := genRPCInfo()
	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
	rpcinfo.Record(ctx, ri, stats.RPCStart, nil)

	// call with retry policy
	var callTimes int32
	policy := BuildFailurePolicy(NewFailurePolicy())
	ri, ok, err := rc.WithRetryIfNeeded(ctx, &policy, retryCall(&callTimes, ri, false), ri, nil, nil)
	test.Assert(t, err == nil, err)
	test.Assert(t, !ok)
	test.Assert(t, ri.Stats().GetEvent(stats.RPCStart).Status() == stats.StatusInfo)
	test.Assert(t, ri.Stats().GetEvent(stats.RPCFinish).Status() == stats.StatusInfo)
	test.Assert(t, ri.To().Address().String() == "10.20.30.40:8888")
	test.Assert(t, atomic.LoadInt32(&callTimes) == 2)
}

var retryWithTransError = func(callTimes, transErrCode int32) RPCCallFunc {
	// fails for the first call if callTimes is initialized to 0
	return func(ctx context.Context, r Retryer, request, response interface{}) (rpcinfo.RPCInfo, error) {
		if atomic.AddInt32(&callTimes, 1) == 1 {
			// first call retry TransErr with specified errCode
			return genRPCInfo(), remote.NewTransErrorWithMsg(transErrCode, "mock")
		} else {
			return genRPCInfoWithRemoteTag(remoteTags), nil
		}
	}
}

func TestNewRetryer4FailureRetry(t *testing.T) {
	t.Run("normal failure retry", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0,
"failure_policy":{"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}}}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, r.(*failureRetryer).enable)
	})

	t.Run("type is failure retry but policy is empty", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0, "failure_policy":{}}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, r.(*failureRetryer).enable)
	})

	t.Run("type is failure retry but policy is mixed", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0,
"mixed_policy":{"retry_delay_ms":10,"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}}}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err != nil, err)
		test.Assert(t, r == nil)
	})

	t.Run("type is failure retry but no policy", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err != nil, err)
		test.Assert(t, r == nil)
	})

	t.Run("type is backup but policy is failure retry", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":1,
"failure_policy":{"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}}}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err != nil, err)
		test.Assert(t, r == nil)
	})

	t.Run("type is failure retry but has multi-policies", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0,
"backup_policy":{"retry_delay_ms":10,"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}},
"failure_policy":{"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}}
}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, r.(*failureRetryer).enable)
	})

	t.Run("type is failure retry but policy json has retry_delay_ms", func(t *testing.T) {
		jsonRet := `{"enable":true,"type":0,
"failure_policy":{"stop_policy":{"max_retry_times":2,"max_duration_ms":0,"disable_chain_stop":false,"ddl_stop":false,"cb_policy":{"error_rate":0.1,"min_sample":200}},"backoff_policy":{"backoff_type":"none"},"retry_same_node":false,"extra":"{\"not_retry_for_timeout\":false,\"rpc_retry_code\":{\"all_error_code\":false,\"error_codes\":[103,1204]},\"biz_retry_code\":{\"all_error_code\":false,\"error_codes\":[]}}","extra_struct":{"not_retry_for_timeout":false,"rpc_retry_code":{"all_error_code":false,"error_codes":[103,1204]},"biz_retry_code":{"all_error_code":false,"error_codes":[]}}}
}`
		var p Policy
		err := sonici.UnmarshalFromString(jsonRet, &p)
		test.Assert(t, err == nil, err)

		r, err := NewRetryer(p, nil, nil)
		test.Assert(t, err == nil, err)
		test.Assert(t, r.(*failureRetryer).enable)
	})
}
