github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/go-grpc-middleware/logging/logrus/shared_test.go (about)

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