github.com/msales/pkg/v3@v3.24.0/grpcx/stats_test.go (about)

     1  package grpcx_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/msales/pkg/v3/grpcx"
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/mock"
    10  	"google.golang.org/grpc/stats"
    11  )
    12  
    13  func TestWithHandlers(t *testing.T) {
    14  	h := grpcx.WithHandlers()
    15  
    16  	assert.Implements(t, (*stats.Handler)(nil), h)
    17  }
    18  
    19  func TestWithRPCStats(t *testing.T) {
    20  	h := grpcx.WithRPCStats(nil)
    21  
    22  	assert.Implements(t, (*stats.Handler)(nil), h)
    23  }
    24  
    25  func TestAggregateHandler_TagRPC(t *testing.T) {
    26  	h1, h2 := new(mockHandler), new(mockHandler)
    27  	ctx := context.Background()
    28  	info := &stats.RPCTagInfo{}
    29  
    30  	h1.On("TagRPC", ctx, info).Return(ctx)
    31  	h2.On("TagRPC", ctx, info).Return(ctx)
    32  
    33  	h := grpcx.WithHandlers(h1, h2)
    34  
    35  	retCtx := h.TagRPC(ctx, info)
    36  
    37  	h1.AssertExpectations(t)
    38  	h2.AssertExpectations(t)
    39  	assert.Equal(t, ctx, retCtx)
    40  }
    41  
    42  func TestAggregateHandler_HandleRPC(t *testing.T) {
    43  	h1, h2 := new(mockHandler), new(mockHandler)
    44  	ctx := context.Background()
    45  	s := &stats.Begin{}
    46  
    47  	h1.On("HandleRPC", ctx, s).Return(ctx)
    48  	h2.On("HandleRPC", ctx, s).Return(ctx)
    49  
    50  	h := grpcx.WithHandlers(h1, h2)
    51  
    52  	h.HandleRPC(ctx, s)
    53  
    54  	h1.AssertExpectations(t)
    55  	h2.AssertExpectations(t)
    56  }
    57  
    58  func TestAggregateHandler_TagConn(t *testing.T) {
    59  	h1, h2 := new(mockHandler), new(mockHandler)
    60  	ctx := context.Background()
    61  	info := &stats.ConnTagInfo{}
    62  
    63  	h1.On("TagConn", ctx, info).Return(ctx)
    64  	h2.On("TagConn", ctx, info).Return(ctx)
    65  
    66  	h := grpcx.WithHandlers(h1, h2)
    67  
    68  	retCtx := h.TagConn(ctx, info)
    69  
    70  	h1.AssertExpectations(t)
    71  	h2.AssertExpectations(t)
    72  	assert.Equal(t, ctx, retCtx)
    73  }
    74  
    75  func TestAggregateHandler_HandleConn(t *testing.T) {
    76  	h1, h2 := new(mockHandler), new(mockHandler)
    77  	ctx := context.Background()
    78  	s := &stats.ConnBegin{}
    79  
    80  	h1.On("HandleConn", ctx, s).Return(ctx)
    81  	h2.On("HandleConn", ctx, s).Return(ctx)
    82  
    83  	h := grpcx.WithHandlers(h1, h2)
    84  
    85  	h.HandleConn(ctx, s)
    86  
    87  	h1.AssertExpectations(t)
    88  	h2.AssertExpectations(t)
    89  }
    90  
    91  type mockHandler struct {
    92  	mock.Mock
    93  }
    94  
    95  func (h *mockHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
    96  	args := h.Called(ctx, info)
    97  
    98  	return args.Get(0).(context.Context)
    99  }
   100  
   101  func (h *mockHandler) HandleRPC(ctx context.Context, s stats.RPCStats) {
   102  	h.Called(ctx, s)
   103  }
   104  
   105  func (h *mockHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
   106  	args := h.Called(ctx, info)
   107  
   108  	return args.Get(0).(context.Context)
   109  }
   110  
   111  func (h *mockHandler) HandleConn(ctx context.Context, s stats.ConnStats) {
   112  	h.Called(ctx, s)
   113  }