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

     1  package grpcx
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/mock"
    11  	"google.golang.org/grpc/stats"
    12  )
    13  
    14  func TestRPCStatsHandler_TagRPC(t *testing.T) {
    15  	h := WithRPCStats(nil)
    16  	info := &stats.RPCTagInfo{FullMethodName: "TestMethod"}
    17  
    18  	ctx := h.TagRPC(context.Background(), info)
    19  
    20  	assert.Equal(t, "TestMethod", ctx.Value(methodKey))
    21  }
    22  
    23  func TestRPCStatsHandler_HandleRPC_End(t *testing.T) {
    24  	ctx := context.WithValue(context.Background(), methodKey, "TestMethod")
    25  	s := new(mockStats)
    26  	s.On("Inc", "rpc.end", int64(1), float32(1.0), []interface{}{"method", "TestMethod", "status", "success"}).Return(nil).Once()
    27  	s.On("Timing", "rpc.time", time.Second, float32(1.0), []interface{}{"method", "TestMethod", "status", "success"}).Return(nil).Once()
    28  	now := time.Now()
    29  
    30  	h := WithRPCStats(s)
    31  
    32  	h.HandleRPC(ctx, &stats.End{BeginTime: now, EndTime: now.Add(time.Second)})
    33  
    34  	s.AssertExpectations(t)
    35  }
    36  
    37  func TestRPCStatsHandler_HandleRPC_End_WithError(t *testing.T) {
    38  	ctx := context.WithValue(context.Background(), methodKey, "TestMethod")
    39  	s := new(mockStats)
    40  	s.On("Inc", "rpc.end", int64(1), float32(1.0), []interface{}{"method", "TestMethod", "status", "error"}).Return(nil).Once()
    41  	s.On("Timing", "rpc.time", time.Second, float32(1.0), []interface{}{"method", "TestMethod", "status", "error"}).Return(nil).Once()
    42  	now := time.Now()
    43  
    44  	h := WithRPCStats(s)
    45  
    46  	h.HandleRPC(ctx, &stats.End{BeginTime: now, EndTime: now.Add(time.Second), Error: errors.New("test: error")})
    47  
    48  	s.AssertExpectations(t)
    49  }
    50  
    51  func TestRPCStatsHandler_HandleRPC_OtherEvent(t *testing.T) {
    52  	ctx := context.WithValue(context.Background(), methodKey, "TestMethod")
    53  	s := new(mockStats)
    54  
    55  	h := WithRPCStats(s)
    56  
    57  	h.HandleRPC(ctx, &stats.InPayload{})
    58  }
    59  
    60  func TestHandler_TagRPC(t *testing.T) {
    61  	h := &handler{}
    62  	ctx := context.Background()
    63  	info := &stats.RPCTagInfo{}
    64  
    65  	retCtx := h.TagRPC(ctx, info)
    66  
    67  	assert.Equal(t, ctx, retCtx)
    68  	assert.Equal(t, &stats.RPCTagInfo{}, info)
    69  }
    70  
    71  func TestHandler_TagConn(t *testing.T) {
    72  	h := &handler{}
    73  	ctx := context.Background()
    74  	info := &stats.ConnTagInfo{}
    75  
    76  	retCtx := h.TagConn(ctx, info)
    77  
    78  	assert.Equal(t, ctx, retCtx)
    79  	assert.Equal(t, &stats.ConnTagInfo{}, info)
    80  }
    81  
    82  type mockStats struct {
    83  	mock.Mock
    84  }
    85  
    86  func (m *mockStats) Inc(name string, value int64, rate float32, tags ...interface{}) error {
    87  	args := m.Called(name, value, rate, tags)
    88  	return args.Error(0)
    89  }
    90  
    91  func (m *mockStats) Dec(name string, value int64, rate float32, tags ...interface{}) error {
    92  	args := m.Called(name, value, rate, tags)
    93  	return args.Error(0)
    94  }
    95  
    96  func (m *mockStats) Gauge(name string, value float64, rate float32, tags ...interface{}) error {
    97  	args := m.Called(name, value, rate, tags)
    98  	return args.Error(0)
    99  }
   100  
   101  func (m *mockStats) Timing(name string, value time.Duration, rate float32, tags ...interface{}) error {
   102  	args := m.Called(name, value, rate, tags)
   103  	return args.Error(0)
   104  }
   105  
   106  func (m *mockStats) Close() error {
   107  	args := m.Called()
   108  	return args.Error(0)
   109  }