gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/go-grpc-middleware/logging/logrus/shared_test.go (about)

     1  package grpc_logrus_test
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"gitee.com/ks-custle/core-gm/go-grpc-middleware/logging/logrus/ctxlogrus"
     7  	"io"
     8  	"testing"
     9  
    10  	grpclogrus "gitee.com/ks-custle/core-gm/go-grpc-middleware/logging/logrus"
    11  	grpcctxtags "gitee.com/ks-custle/core-gm/go-grpc-middleware/tags"
    12  	grpctesting "gitee.com/ks-custle/core-gm/go-grpc-middleware/testing"
    13  	pbtestproto "gitee.com/ks-custle/core-gm/go-grpc-middleware/testing/testproto"
    14  	"gitee.com/ks-custle/core-gm/grpc/codes"
    15  	"gitee.com/ks-custle/core-gm/net/context"
    16  	"github.com/sirupsen/logrus"
    17  )
    18  
    19  var (
    20  	goodPing = &pbtestproto.PingRequest{Value: "something", SleepTimeMs: 9999}
    21  )
    22  
    23  type loggingPingService struct {
    24  	pbtestproto.TestServiceServer
    25  }
    26  
    27  func customCodeToLevel(c codes.Code) logrus.Level {
    28  	if c == codes.Unauthenticated {
    29  		// Make this a special case for tests, and an error.
    30  		return logrus.ErrorLevel
    31  	}
    32  	level := grpclogrus.DefaultCodeToLevel(c)
    33  	return level
    34  }
    35  
    36  func (s *loggingPingService) Ping(ctx context.Context, ping *pbtestproto.PingRequest) (*pbtestproto.PingResponse, error) {
    37  	grpcctxtags.Extract(ctx).Set("custom_tags.string", "something").Set("custom_tags.int", 1337)
    38  	// ctx_logrus.AddFields is deprecated, use the ctxlogrus.Extract instead.
    39  	//ctx_logrus.AddFields(ctx, logrus.Fields{"custom_field": "custom_value"})
    40  	ctxlogrus.AddFields(ctx, logrus.Fields{"custom_field": "custom_value"})
    41  	// ctx_logrus.Extract is deprecated, use the ctxlogrus.Extract instead.
    42  	//ctx_logrus.Extract(ctx).Info("some ping")
    43  	ctxlogrus.Extract(ctx).Info("some ping")
    44  	return s.TestServiceServer.Ping(ctx, ping)
    45  }
    46  
    47  func (s *loggingPingService) PingError(ctx context.Context, ping *pbtestproto.PingRequest) (*pbtestproto.Empty, error) {
    48  	return s.TestServiceServer.PingError(ctx, ping)
    49  }
    50  
    51  func (s *loggingPingService) PingList(ping *pbtestproto.PingRequest, stream pbtestproto.TestService_PingListServer) error {
    52  	grpcctxtags.Extract(stream.Context()).Set("custom_tags.string", "something").Set("custom_tags.int", 1337)
    53  	// ctx_logrus.AddFields is deprecated, use the ctxlogrus.Extract instead.
    54  	//ctx_logrus.AddFields(stream.Context(), logrus.Fields{"custom_field": "custom_value"})
    55  	ctxlogrus.AddFields(stream.Context(), logrus.Fields{"custom_field": "custom_value"})
    56  	// ctx_logrus.Extract is deprecated, use the ctxlogrus.Extract instead.
    57  	//ctx_logrus.Extract(stream.Context()).Info("some pinglist")
    58  	ctxlogrus.Extract(stream.Context()).Info("some pinglist")
    59  	return s.TestServiceServer.PingList(ping, stream)
    60  }
    61  
    62  func (s *loggingPingService) PingEmpty(ctx context.Context, empty *pbtestproto.Empty) (*pbtestproto.PingResponse, error) {
    63  	return s.TestServiceServer.PingEmpty(ctx, empty)
    64  }
    65  
    66  type logrusBaseSuite struct {
    67  	*grpctesting.InterceptorTestSuite
    68  	mutexBuffer *grpctesting.MutexReadWriter
    69  	buffer      *bytes.Buffer
    70  	logger      *logrus.Logger
    71  }
    72  
    73  func newLogrusBaseSuite(t *testing.T) *logrusBaseSuite {
    74  	b := &bytes.Buffer{}
    75  	muB := grpctesting.NewMutexReadWriter(b)
    76  	logger := logrus.New()
    77  	logger.Out = muB
    78  	logger.Formatter = &logrus.JSONFormatter{DisableTimestamp: true}
    79  	return &logrusBaseSuite{
    80  		logger:      logger,
    81  		buffer:      b,
    82  		mutexBuffer: muB,
    83  		InterceptorTestSuite: &grpctesting.InterceptorTestSuite{
    84  			TestService: &loggingPingService{&grpctesting.TestPingService{T: t}},
    85  		},
    86  	}
    87  }
    88  
    89  func (s *logrusBaseSuite) SetupTest() {
    90  	s.mutexBuffer.Lock()
    91  	s.buffer.Reset()
    92  	s.mutexBuffer.Unlock()
    93  }
    94  
    95  func (s *logrusBaseSuite) getOutputJSONs() []map[string]interface{} {
    96  	ret := make([]map[string]interface{}, 0)
    97  	dec := json.NewDecoder(s.mutexBuffer)
    98  
    99  	for {
   100  		var val map[string]interface{}
   101  		err := dec.Decode(&val)
   102  		if err == io.EOF {
   103  			break
   104  		}
   105  		if err != nil {
   106  			s.T().Fatalf("failed decoding output from Logrus JSON: %v", err)
   107  		}
   108  
   109  		ret = append(ret, val)
   110  	}
   111  
   112  	return ret
   113  }