github.com/pinpoint-apm/pinpoint-go-agent@v1.4.1-0.20240110120318-a50c2eb18c8c/mock_grpc.go (about)

     1  package pinpoint
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/golang/mock/gomock"
    10  	lru "github.com/hashicorp/golang-lru"
    11  	pb "github.com/pinpoint-apm/pinpoint-go-agent/protobuf"
    12  	"github.com/pinpoint-apm/pinpoint-go-agent/protobuf/mock"
    13  )
    14  
    15  func newTestAgent(config *Config) *agent {
    16  	a := &agent{
    17  		appName:   "testApp",
    18  		agentID:   "testAgent",
    19  		appType:   ServiceTypeGoApp,
    20  		startTime: time.Now().UnixNano() / int64(time.Millisecond),
    21  		enable:    true,
    22  		spanChan:  make(chan *span, cacheSize),
    23  		metaChan:  make(chan interface{}, cacheSize),
    24  		sampler:   newBasicTraceSampler(newRateSampler(1)),
    25  		config:    config,
    26  	}
    27  	a.errorCache, _ = lru.New(cacheSize)
    28  	a.sqlCache, _ = lru.New(cacheSize)
    29  	a.sqlUidCache, _ = lru.New(cacheSize)
    30  	a.apiCache, _ = lru.New(cacheSize)
    31  	a.config.offGrpc = true
    32  
    33  	return a
    34  }
    35  
    36  //mock grpc
    37  
    38  type mockAgentGrpcClient struct {
    39  	client *mock.MockAgentClient
    40  	stream *mock.MockAgent_PingSessionClient
    41  }
    42  
    43  func (agentGrpcClient *mockAgentGrpcClient) RequestAgentInfo(ctx context.Context, agentinfo *pb.PAgentInfo) (*pb.PResult, error) {
    44  	_ = agentGrpcClient.client.EXPECT().RequestAgentInfo(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    45  	return agentGrpcClient.client.RequestAgentInfo(ctx, agentinfo)
    46  }
    47  
    48  func (agentGrpcClient *mockAgentGrpcClient) PingSession(ctx context.Context) (pb.Agent_PingSessionClient, error) {
    49  	_ = agentGrpcClient.client.EXPECT().PingSession(gomock.Any()).Return(agentGrpcClient.stream, nil)
    50  	return agentGrpcClient.client.PingSession(ctx)
    51  }
    52  
    53  type mockMetaGrpcClient struct {
    54  	client *mock.MockMetadataClient
    55  }
    56  
    57  func (metaGrpcClient *mockMetaGrpcClient) RequestApiMetaData(ctx context.Context, in *pb.PApiMetaData) (*pb.PResult, error) {
    58  	_ = metaGrpcClient.client.EXPECT().RequestApiMetaData(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    59  	return metaGrpcClient.client.RequestApiMetaData(ctx, in)
    60  }
    61  
    62  func (metaGrpcClient *mockMetaGrpcClient) RequestSqlMetaData(ctx context.Context, in *pb.PSqlMetaData) (*pb.PResult, error) {
    63  	_ = metaGrpcClient.client.EXPECT().RequestSqlMetaData(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    64  	return metaGrpcClient.client.RequestSqlMetaData(ctx, in)
    65  }
    66  
    67  func (metaGrpcClient *mockMetaGrpcClient) RequestSqlUidMetaData(ctx context.Context, in *pb.PSqlUidMetaData) (*pb.PResult, error) {
    68  	//_ = metaGrpcClient.client.EXPECT().RequestSqlMetaData(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    69  	//return metaGrpcClient.client.RequestSqlUidMetaData(ctx, in)
    70  	return nil, nil
    71  }
    72  
    73  func (metaGrpcClient *mockMetaGrpcClient) RequestStringMetaData(ctx context.Context, in *pb.PStringMetaData) (*pb.PResult, error) {
    74  	_ = metaGrpcClient.client.EXPECT().RequestStringMetaData(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    75  	return metaGrpcClient.client.RequestStringMetaData(ctx, in)
    76  }
    77  
    78  func (metaGrpcClient *mockMetaGrpcClient) RequestExceptionMetaData(ctx context.Context, in *pb.PExceptionMetaData) (*pb.PResult, error) {
    79  	//_ = metaGrpcClient.client.EXPECT().RequestExceptionMetaData(gomock.Any(), gomock.Any()).Return(&pb.PResult{Success: true, Message: "success"}, nil)
    80  	//return metaGrpcClient.client.RequestExceptionMetaData(ctx, in)
    81  	return nil, nil
    82  }
    83  
    84  func newMockAgentGrpc(agent *agent, t *testing.T) *agentGrpc {
    85  	ctrl := gomock.NewController(t)
    86  	stream := mock.NewMockAgent_PingSessionClient(ctrl)
    87  	agentClient := mockAgentGrpcClient{mock.NewMockAgentClient(ctrl), stream}
    88  	metadataClient := mockMetaGrpcClient{mock.NewMockMetadataClient(ctrl)}
    89  
    90  	return &agentGrpc{nil, &agentClient, &metadataClient, -1, nil, agent}
    91  }
    92  
    93  func newMockAgentGrpcPing(agent *agent, t *testing.T) *agentGrpc {
    94  	ctrl := gomock.NewController(t)
    95  	stream := mock.NewMockAgent_PingSessionClient(ctrl)
    96  	agentClient := mockAgentGrpcClient{mock.NewMockAgentClient(ctrl), stream}
    97  	metadataClient := mockMetaGrpcClient{mock.NewMockMetadataClient(ctrl)}
    98  
    99  	stream.EXPECT().Send(gomock.Any()).Return(nil)
   100  	stream.EXPECT().Recv().Return(nil, nil)
   101  
   102  	return &agentGrpc{nil, &agentClient, &metadataClient, -1, nil, agent}
   103  }
   104  
   105  type mockSpanGrpcClient struct {
   106  	client *mock.MockSpanClient
   107  	stream *mock.MockSpan_SendSpanClient
   108  }
   109  
   110  func (spanGrpcClient *mockSpanGrpcClient) SendSpan(ctx context.Context) (pb.Span_SendSpanClient, error) {
   111  	_ = spanGrpcClient.client.EXPECT().SendSpan(gomock.Any()).Return(spanGrpcClient.stream, nil)
   112  	return spanGrpcClient.client.SendSpan(ctx)
   113  }
   114  
   115  func newMockSpanGrpc(agent *agent, t *testing.T) *spanGrpc {
   116  	ctrl := gomock.NewController(t)
   117  	stream := mock.NewMockSpan_SendSpanClient(ctrl)
   118  	spanClient := mockSpanGrpcClient{mock.NewMockSpanClient(ctrl), stream}
   119  
   120  	stream.EXPECT().Send(gomock.Any()).Return(nil)
   121  
   122  	return &spanGrpc{nil, &spanClient, nil, agent}
   123  }
   124  
   125  type mockStaGrpcClient struct {
   126  	client *mock.MockStatClient
   127  	stream *mock.MockStat_SendAgentStatClient
   128  }
   129  
   130  func (statGrpcClient *mockStaGrpcClient) SendAgentStat(ctx context.Context) (pb.Stat_SendAgentStatClient, error) {
   131  	_ = statGrpcClient.client.EXPECT().SendAgentStat(gomock.Any()).Return(statGrpcClient.stream, nil)
   132  	return statGrpcClient.client.SendAgentStat(ctx)
   133  }
   134  
   135  func newMockStatGrpc(agent *agent, t *testing.T) *statGrpc {
   136  	ctrl := gomock.NewController(t)
   137  	stream := mock.NewMockStat_SendAgentStatClient(ctrl)
   138  	statClient := mockStaGrpcClient{mock.NewMockStatClient(ctrl), stream}
   139  
   140  	stream.EXPECT().Send(gomock.Any()).Return(nil)
   141  
   142  	return &statGrpc{nil, &statClient, nil, agent}
   143  }
   144  
   145  type mockRetryStaGrpcClient struct {
   146  	client *mock.MockStatClient
   147  	stream *mock.MockStat_SendAgentStatClient
   148  	retry  int
   149  }
   150  
   151  func (statGrpcClient *mockRetryStaGrpcClient) SendAgentStat(ctx context.Context) (pb.Stat_SendAgentStatClient, error) {
   152  	if statGrpcClient.retry < 3 {
   153  		time.Sleep(1 * time.Second)
   154  		_ = statGrpcClient.client.EXPECT().SendAgentStat(gomock.Any()).Return(nil, errors.New(""))
   155  	} else {
   156  		_ = statGrpcClient.client.EXPECT().SendAgentStat(gomock.Any()).Return(statGrpcClient.stream, nil)
   157  	}
   158  	statGrpcClient.retry++
   159  	return statGrpcClient.client.SendAgentStat(ctx)
   160  }
   161  
   162  func newRetryMockStatGrpc(agent *agent, t *testing.T) *statGrpc {
   163  	ctrl := gomock.NewController(t)
   164  	stream := mock.NewMockStat_SendAgentStatClient(ctrl)
   165  	statClient := mockRetryStaGrpcClient{mock.NewMockStatClient(ctrl), stream, 0}
   166  
   167  	stream.EXPECT().Send(gomock.Any()).Return(nil)
   168  	//stream.EXPECT().Send(gomock.Any()).Return(errors.New("stat send fail"))
   169  
   170  	return &statGrpc{nil, &statClient, nil, agent}
   171  }