1 // Copyright (c) 2017 Uber Technologies, Inc. 2 // 3 // Permission is hereby granted, free of charge, to any person obtaining a copy 4 // of this software and associated documentation files (the "Software"), to deal 5 // in the Software without restriction, including without limitation the rights 6 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 // copies of the Software, and to permit persons to whom the Software is 8 // furnished to do so, subject to the following conditions: 9 // 10 // The above copyright notice and this permission notice shall be included in 11 // all copies or substantial portions of the Software. 12 // 13 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 19 // THE SOFTWARE. 20 21 package zaptest 22 23 import ( 24 "bytes" 25 26 "go.uber.org/zap" 27 "go.uber.org/zap/zapcore" 28 ) 29 30 // LoggerOption configures the test logger built by NewLogger. 31 type LoggerOption interface { 32 applyLoggerOption(*loggerOptions) 33 } 34 35 type loggerOptions struct { 36 Level zapcore.LevelEnabler 37 zapOptions []zap.Option 38 } 39 40 type loggerOptionFunc func(*loggerOptions) 41 42 func (f loggerOptionFunc) applyLoggerOption(opts *loggerOptions) { 43 f(opts) 44 } 45 46 // Level controls which messages are logged by a test Logger built by 47 // NewLogger. 48 func Level(enab zapcore.LevelEnabler) LoggerOption { 49 return loggerOptionFunc(func(opts *loggerOptions) { 50 opts.Level = enab 51 }) 52 } 53 54 // WrapOptions adds zap.Option's to a test Logger built by NewLogger. 55 func WrapOptions(zapOpts ...zap.Option) LoggerOption { 56 return loggerOptionFunc(func(opts *loggerOptions) { 57 opts.zapOptions = zapOpts 58 }) 59 } 60 61 // NewLogger builds a new Logger that logs all messages to the given 62 // testing.TB. 63 // 64 // logger := zaptest.NewLogger(t) 65 // 66 // Use this with a *testing.T or *testing.B to get logs which get printed only 67 // if a test fails or if you ran go test -v. 68 // 69 // The returned logger defaults to logging debug level messages and above. 70 // This may be changed by passing a zaptest.Level during construction. 71 // 72 // logger := zaptest.NewLogger(t, zaptest.Level(zap.WarnLevel)) 73 // 74 // You may also pass zap.Option's to customize test logger. 75 // 76 // logger := zaptest.NewLogger(t, zaptest.WrapOptions(zap.AddCaller())) 77 func NewLogger(t TestingT, opts ...LoggerOption) *zap.Logger { 78 cfg := loggerOptions{ 79 Level: zapcore.DebugLevel, 80 } 81 for _, o := range opts { 82 o.applyLoggerOption(&cfg) 83 } 84 85 writer := NewTestingWriter(t) 86 zapOptions := []zap.Option{ 87 // Send zap errors to the same writer and mark the test as failed if 88 // that happens. 89 zap.ErrorOutput(writer.WithMarkFailed(true)), 90 } 91 zapOptions = append(zapOptions, cfg.zapOptions...) 92 93 return zap.New( 94 zapcore.NewCore( 95 zapcore.NewConsoleEncoder(zap.NewDevelopmentEncoderConfig()), 96 writer, 97 cfg.Level, 98 ), 99 zapOptions..., 100 ) 101 } 102 103 // TestingWriter is a WriteSyncer that writes to the given testing.TB. 104 type TestingWriter struct { 105 t TestingT 106 107 // If true, the test will be marked as failed if this TestingWriter is 108 // ever used. 109 markFailed bool 110 } 111 112 // NewTestingWriter builds a new TestingWriter that writes to the given 113 // testing.TB. 114 // 115 // Use this if you need more flexibility when creating *zap.Logger 116 // than zaptest.NewLogger() provides. 117 // 118 // E.g., if you want to use custom core with zaptest.TestingWriter: 119 // 120 // encoder := newCustomEncoder() 121 // writer := zaptest.NewTestingWriter(t) 122 // level := zap.NewAtomicLevelAt(zapcore.DebugLevel) 123 // 124 // core := newCustomCore(encoder, writer, level) 125 // 126 // logger := zap.New(core, zap.AddCaller()) 127 func NewTestingWriter(t TestingT) TestingWriter { 128 return TestingWriter{t: t} 129 } 130 131 // WithMarkFailed returns a copy of this TestingWriter with markFailed set to 132 // the provided value. 133 func (w TestingWriter) WithMarkFailed(v bool) TestingWriter { 134 w.markFailed = v 135 return w 136 } 137 138 // Write writes bytes from p to the underlying testing.TB. 139 func (w TestingWriter) Write(p []byte) (n int, err error) { 140 n = len(p) 141 142 // Strip trailing newline because t.Log always adds one. 143 p = bytes.TrimRight(p, "\n") 144 145 // Note: t.Log is safe for concurrent use. 146 w.t.Logf("%s", p) 147 if w.markFailed { 148 w.t.Fail() 149 } 150 151 return n, nil 152 } 153 154 // Sync commits the current contents (a no-op for TestingWriter). 155 func (w TestingWriter) Sync() error { 156 return nil 157 } 158