github.com/cloudwego/kitex@v0.9.0/pkg/circuitbreak/circuitbreak_test.go

/*
 * 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 circuitbreak

import (
	"context"
	"errors"
	"testing"

	"github.com/bytedance/gopkg/cloud/circuitbreaker"

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

type ctxKeyType int

const (
	ctxKey ctxKeyType = iota
	ctxEnabled
	ctxErrorType
)

var (
	errFake     = errors.New("fake error")
	errDecorate = errors.New("decorate error")
	allowed     = "Allowed"
)

func invoke(ctx context.Context, request, response interface{}) error {
	return errFake
}

func mockGetKey(ctx context.Context, request interface{}) (key string, enabled bool) {
	key = ctx.Value(ctxKey).(string)
	enabled = ctx.Value(ctxEnabled).(bool)
	return
}

func mockGetErrorType(ctx context.Context, request, response interface{}, err error) ErrorType {
	t, ok := ctx.Value(ctxErrorType).(ErrorType)
	if !ok {
		return TypeIgnorable
	}
	return t
}

func mockDecorateError(ctx context.Context, request interface{}, err error) error {
	return errDecorate
}

type mockPanel struct {
	circuitbreaker.Panel
	judged          bool
	timeoutRecorded bool
	failRecorded    bool
	succeedRecorded bool
}

func (m *mockPanel) IsAllowed(key string) bool {
	m.judged = true
	return key == allowed
}

func (m *mockPanel) Timeout(key string) {
	m.timeoutRecorded = true
}

func (m *mockPanel) Fail(key string) {
	m.failRecorded = true
}

func (m *mockPanel) Succeed(key string) {
	m.succeedRecorded = true
}

func TestNewCircuitBreakerMW(t *testing.T) {
	ctl := Control{
		GetKey:        mockGetKey,
		GetErrorType:  mockGetErrorType,
		DecorateError: mockDecorateError,
	}
	panel, err := circuitbreaker.NewPanel(nil, circuitbreaker.Options{})
	if err != nil {
		t.Fatal(err)
	}
	mp := &mockPanel{Panel: panel}
	cbMW := NewCircuitBreakerMW(ctl, mp)
	// test disabled
	ctx := context.Background()
	ctx = context.WithValue(ctx, ctxKey, allowed)
	ctx = context.WithValue(ctx, ctxEnabled, false)
	test.Assert(t, errors.Is(cbMW(invoke)(ctx, nil, nil), errFake))
	test.Assert(t, !mp.judged)
	// test enabled and allowed
	ctx = context.WithValue(ctx, ctxKey, allowed)
	ctx = context.WithValue(ctx, ctxEnabled, true)
	test.Assert(t, errors.Is(cbMW(invoke)(ctx, nil, nil), errFake))
	test.Assert(t, mp.judged)
	// test enabled and not allowed
	ctx = context.WithValue(ctx, ctxKey, "you should not pass")
	ctx = context.WithValue(ctx, ctxEnabled, true)
	test.Assert(t, errors.Is(cbMW(invoke)(ctx, nil, nil), errDecorate))
	test.Assert(t, mp.judged)
}

func TestRecordStat(t *testing.T) {
	ctl := Control{
		GetKey:        mockGetKey,
		GetErrorType:  mockGetErrorType,
		DecorateError: mockDecorateError,
	}
	panel, err := circuitbreaker.NewPanel(nil, circuitbreaker.Options{})
	if err != nil {
		t.Fatal(err)
	}
	// test timeout
	mp := &mockPanel{Panel: panel}
	cbMW := NewCircuitBreakerMW(ctl, mp)
	ctx := context.Background()
	ctx = context.WithValue(ctx, ctxKey, allowed)
	ctx = context.WithValue(ctx, ctxEnabled, true)
	ctx = context.WithValue(ctx, ctxErrorType, TypeTimeout)
	cbMW(invoke)(ctx, nil, nil)
	test.Assert(t, mp.timeoutRecorded)
	// test failure
	mp = &mockPanel{Panel: panel}
	cbMW = NewCircuitBreakerMW(ctl, mp)
	ctx = context.Background()
	ctx = context.WithValue(ctx, ctxKey, allowed)
	ctx = context.WithValue(ctx, ctxEnabled, true)
	ctx = context.WithValue(ctx, ctxErrorType, TypeFailure)
	cbMW(invoke)(ctx, nil, nil)
	test.Assert(t, mp.failRecorded)
	// test success
	mp = &mockPanel{Panel: panel}
	cbMW = NewCircuitBreakerMW(ctl, mp)
	ctx = context.Background()
	ctx = context.WithValue(ctx, ctxKey, allowed)
	ctx = context.WithValue(ctx, ctxEnabled, true)
	ctx = context.WithValue(ctx, ctxErrorType, TypeSuccess)
	cbMW(invoke)(ctx, nil, nil)
	test.Assert(t, mp.succeedRecorded)
}

func TestErrorType(t *testing.T) {
	err1 := kerrors.ErrRPCTimeout
	wrapErr := WrapErrorWithType(err1, TypeIgnorable)
	test.Assert(t, errors.Is(wrapErr, kerrors.ErrRPCTimeout))
	test.Assert(t, errors.Unwrap(wrapErr) == err1)
	test.Assert(t, wrapErr.TypeForCircuitBreaker() == TypeIgnorable)

	err2 := kerrors.ErrRemoteOrNetwork.WithCause(errors.New("mock"))
	wrapErr = WrapErrorWithType(err2, TypeIgnorable)
	test.Assert(t, errors.Is(wrapErr, kerrors.ErrRemoteOrNetwork))
	test.Assert(t, errors.Is(errors.Unwrap(wrapErr), kerrors.ErrRemoteOrNetwork))
	test.Assert(t, wrapErr.TypeForCircuitBreaker() == TypeIgnorable)

	err3 := kerrors.ErrRemoteOrNetwork.WithCause(wrapErr)
	var e *errorWrapperWithType
	ok := errors.As(err3, &e)
	test.Assert(t, ok)
	test.Assert(t, e.TypeForCircuitBreaker() == TypeIgnorable)
}