github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/serverinterceptors/sheddinginterceptor_test.go (about)

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/lingyao2333/mo-zero/core/load"
     8  	"github.com/lingyao2333/mo-zero/core/stat"
     9  	"github.com/stretchr/testify/assert"
    10  	"google.golang.org/grpc"
    11  )
    12  
    13  func TestUnarySheddingInterceptor(t *testing.T) {
    14  	tests := []struct {
    15  		name      string
    16  		allow     bool
    17  		handleErr error
    18  		expect    error
    19  	}{
    20  		{
    21  			name:      "allow",
    22  			allow:     true,
    23  			handleErr: nil,
    24  			expect:    nil,
    25  		},
    26  		{
    27  			name:      "allow",
    28  			allow:     true,
    29  			handleErr: context.DeadlineExceeded,
    30  			expect:    context.DeadlineExceeded,
    31  		},
    32  		{
    33  			name:      "reject",
    34  			allow:     false,
    35  			handleErr: nil,
    36  			expect:    load.ErrServiceOverloaded,
    37  		},
    38  	}
    39  
    40  	for _, test := range tests {
    41  		test := test
    42  		t.Run(test.name, func(t *testing.T) {
    43  			t.Parallel()
    44  
    45  			shedder := mockedShedder{allow: test.allow}
    46  			metrics := stat.NewMetrics("mock")
    47  			interceptor := UnarySheddingInterceptor(shedder, metrics)
    48  			_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
    49  				FullMethod: "/",
    50  			}, func(ctx context.Context, req interface{}) (interface{}, error) {
    51  				return nil, test.handleErr
    52  			})
    53  			assert.Equal(t, test.expect, err)
    54  		})
    55  	}
    56  }
    57  
    58  type mockedShedder struct {
    59  	allow bool
    60  }
    61  
    62  func (m mockedShedder) Allow() (load.Promise, error) {
    63  	if m.allow {
    64  		return mockedPromise{}, nil
    65  	}
    66  
    67  	return nil, load.ErrServiceOverloaded
    68  }
    69  
    70  type mockedPromise struct{}
    71  
    72  func (m mockedPromise) Pass() {
    73  }
    74  
    75  func (m mockedPromise) Fail() {
    76  }