1 package grpc_logrus_test
2
3 import (
4 "bytes"
5 "context"
6 "encoding/json"
7 "io"
8 "testing"
9
10 grpc_logrus "github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus"
11 "github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
12 grpc_ctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
13 grpc_testing "github.com/grpc-ecosystem/go-grpc-middleware/testing"
14 pb_testproto "github.com/grpc-ecosystem/go-grpc-middleware/testing/testproto"
15 "github.com/sirupsen/logrus"
16 "google.golang.org/grpc/codes"
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
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 ctxlogrus.AddFields(ctx, logrus.Fields{"custom_field": "custom_value"})
39 ctxlogrus.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 ctxlogrus.AddFields(stream.Context(), logrus.Fields{"custom_field": "custom_value"})
50 ctxlogrus.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 timestampFormat string
64 }
65
66 func newLogrusBaseSuite(t *testing.T) *logrusBaseSuite {
67 b := &bytes.Buffer{}
68 muB := grpc_testing.NewMutexReadWriter(b)
69 logger := logrus.New()
70 logger.Out = muB
71 logger.Formatter = &logrus.JSONFormatter{DisableTimestamp: true}
72 return &logrusBaseSuite{
73 logger: logger,
74 buffer: b,
75 mutexBuffer: muB,
76 InterceptorTestSuite: &grpc_testing.InterceptorTestSuite{
77 TestService: &loggingPingService{&grpc_testing.TestPingService{T: t}},
78 },
79 }
80 }
81
82 func (s *logrusBaseSuite) SetupTest() {
83 s.mutexBuffer.Lock()
84 s.buffer.Reset()
85 s.mutexBuffer.Unlock()
86 }
87
88 func (s *logrusBaseSuite) 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 }
107
108 func StubMessageProducer(ctx context.Context, format string, level logrus.Level, code codes.Code, err error, fields logrus.Fields) {
109 if err != nil {
110 fields[logrus.ErrorKey] = err
111 }
112 format = "custom message"
113 entry := ctxlogrus.Extract(ctx).WithContext(ctx).WithFields(fields)
114 switch level {
115 case logrus.DebugLevel:
116 entry.Debugf(format)
117 case logrus.InfoLevel:
118 entry.Infof(format)
119 case logrus.WarnLevel:
120 entry.Warningf(format)
121 case logrus.ErrorLevel:
122 entry.Errorf(format)
123 case logrus.FatalLevel:
124 entry.Fatalf(format)
125 case logrus.PanicLevel:
126 entry.Panicf(format)
127 }
128 }
129
View as plain text