logger.go

package testutils

import (
	"os"
	"testing"

	"github.com/rs/zerolog"
)

// NewTestLogger creates a TestLogger
func NewTestLogger(tb testing.TB) *TestLogger {
	return &TestLogger{tb}
}

// TestLogger logs to a t.Log function
type TestLogger struct {
	tb testing.TB
}

// Logger returns a zerolog Logger
func (tl *TestLogger) Logger() zerolog.Logger {
	level := zerolog.DebugLevel
	if os.Getenv("TEST_TRACE") != "" {
		level = zerolog.TraceLevel
	}
	if os.Getenv("TEST_QUIET") != "" {
		level = zerolog.InfoLevel
	}
	return zerolog.New(zerolog.ConsoleWriter{Out: tl}).
		With().Timestamp().Logger().
		Level(level)
}

// Write writes the given string to t.Logf
func (tl *TestLogger) Write(m []byte) (int, error) {
	tl.tb.Logf(string(m))
	return len(m), nil
}

// SetTB changes the current TB, and returns a function to get back to the
// previous one
func (tl *TestLogger) SetTB(tb testing.TB) func() {
	otb := tl.tb
	tl.tb = tb
	return func() {
		tl.tb = otb
	}
}

// GetLogger returns a test Logger
func GetLogger(tb testing.TB) zerolog.Logger {
	return NewTestLogger(tb).Logger()
}