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

     1  package grpc_zap_test
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"io"
     7  	"testing"
     8  
     9  	grpc_ctxtags "github.com/hxx258456/ccgo/go-grpc-middleware/tags"
    10  	ctx_zap "github.com/hxx258456/ccgo/go-grpc-middleware/tags/zap"
    11  	grpc_testing "github.com/hxx258456/ccgo/go-grpc-middleware/testing"
    12  	pb_testproto "github.com/hxx258456/ccgo/go-grpc-middleware/testing/testproto"
    13  	"github.com/hxx258456/ccgo/net/context"
    14  	"go.uber.org/zap"
    15  	"go.uber.org/zap/zapcore"
    16  )
    17  
    18  var (
    19  	goodPing = &pb_testproto.PingRequest{Value: "something", SleepTimeMs: 9999}
    20  )
    21  
    22  type loggingPingService struct {
    23  	pb_testproto.TestServiceServer
    24  }
    25  
    26  func (s *loggingPingService) Ping(ctx context.Context, ping *pb_testproto.PingRequest) (*pb_testproto.PingResponse, error) {
    27  	grpc_ctxtags.Extract(ctx).Set("custom_tags.string", "something").Set("custom_tags.int", 1337)
    28  	ctx_zap.AddFields(ctx, zap.String("custom_field", "custom_value"))
    29  	ctx_zap.Extract(ctx).Info("some ping")
    30  	return s.TestServiceServer.Ping(ctx, ping)
    31  }
    32  
    33  func (s *loggingPingService) PingError(ctx context.Context, ping *pb_testproto.PingRequest) (*pb_testproto.Empty, error) {
    34  	return s.TestServiceServer.PingError(ctx, ping)
    35  }
    36  
    37  func (s *loggingPingService) PingList(ping *pb_testproto.PingRequest, stream pb_testproto.TestService_PingListServer) error {
    38  	grpc_ctxtags.Extract(stream.Context()).Set("custom_tags.string", "something").Set("custom_tags.int", 1337)
    39  	ctx_zap.Extract(stream.Context()).Info("some pinglist")
    40  	return s.TestServiceServer.PingList(ping, stream)
    41  }
    42  
    43  func (s *loggingPingService) PingEmpty(ctx context.Context, empty *pb_testproto.Empty) (*pb_testproto.PingResponse, error) {
    44  	return s.TestServiceServer.PingEmpty(ctx, empty)
    45  }
    46  
    47  func newBaseZapSuite(t *testing.T) *zapBaseSuite {
    48  	b := &bytes.Buffer{}
    49  	muB := grpc_testing.NewMutexReadWriter(b)
    50  	zap.NewDevelopmentConfig()
    51  	jsonEncoder := zapcore.NewJSONEncoder(zapcore.EncoderConfig{
    52  		TimeKey:        "ts",
    53  		LevelKey:       "level",
    54  		NameKey:        "logger",
    55  		CallerKey:      "caller",
    56  		MessageKey:     "msg",
    57  		StacktraceKey:  "stacktrace",
    58  		EncodeLevel:    zapcore.LowercaseLevelEncoder,
    59  		EncodeTime:     zapcore.EpochTimeEncoder,
    60  		EncodeDuration: zapcore.SecondsDurationEncoder,
    61  	})
    62  	core := zapcore.NewCore(jsonEncoder, zapcore.AddSync(muB), zap.LevelEnablerFunc(func(zapcore.Level) bool { return true }))
    63  	log := zap.New(core)
    64  	s := &zapBaseSuite{
    65  		log:         log,
    66  		buffer:      b,
    67  		mutexBuffer: muB,
    68  		InterceptorTestSuite: &grpc_testing.InterceptorTestSuite{
    69  			TestService: &loggingPingService{&grpc_testing.TestPingService{T: t}},
    70  		},
    71  	}
    72  	return s
    73  }
    74  
    75  type zapBaseSuite struct {
    76  	*grpc_testing.InterceptorTestSuite
    77  	mutexBuffer *grpc_testing.MutexReadWriter
    78  	buffer      *bytes.Buffer
    79  	log         *zap.Logger
    80  }
    81  
    82  func (s *zapBaseSuite) SetupTest() {
    83  	s.mutexBuffer.Lock()
    84  	s.buffer.Reset()
    85  	s.mutexBuffer.Unlock()
    86  }
    87  
    88  func (s *zapBaseSuite) getOutputJSONs() []map[string]interface{} {
    89  	ret := make([]map[string]interface{}, 0)
    90  	dec := json.NewDecoder(s.mutexBuffer)
    91  
    92  	for {
    93  		var val map[string]interface{}
    94  		err := dec.Decode(&val)
    95  		if err == io.EOF {
    96  			break
    97  		}
    98  		if err != nil {
    99  			s.T().Fatalf("failed decoding output from Logrus JSON: %v", err)
   100  		}
   101  
   102  		ret = append(ret, val)
   103  	}
   104  
   105  	return ret
   106  }