testhub/logger/logger.go

package logger

import (
	"log"
	"net/http"
	"time"

	"github.com/sirupsen/logrus"
)

const (
	Disabled = iota
	InfoLevel
	WarningLevel
	DebugLevel
)

type (
	// struct for holding response details
	responseData struct {
		status int
		size   int
	}

	// our http.ResponseWriter implementation
	loggingResponseWriter struct {
		http.ResponseWriter // compose original http.ResponseWriter
		responseData        *responseData
	}
)

var loggerLevel int = InfoLevel

func (r *loggingResponseWriter) Write(b []byte) (int, error) {
	size, err := r.ResponseWriter.Write(b) // write response using original http.ResponseWriter
	r.responseData.size += size            // capture size
	return size, err
}

func (r *loggingResponseWriter) WriteHeader(statusCode int) {
	r.ResponseWriter.WriteHeader(statusCode) // write status code using original http.ResponseWriter
	r.responseData.status = statusCode       // capture status code
}

func WithLogging(h http.Handler) http.Handler {
	loggingFn := func(rw http.ResponseWriter, req *http.Request) {
		start := time.Now()

		responseData := &responseData{
			status: 0,
			size:   0,
		}
		lrw := loggingResponseWriter{
			ResponseWriter: rw, // compose original http.ResponseWriter
			responseData:   responseData,
		}
		h.ServeHTTP(&lrw, req) // inject our implementation of http.ResponseWriter

		duration := time.Since(start)

		logrus.WithFields(logrus.Fields{
			"uri":      req.RequestURI,
			"method":   req.Method,
			"status":   responseData.status,
			"duration": duration,
			"size":     responseData.size,
		}).Info("request completed")
	}
	return http.HandlerFunc(loggingFn)
}

func SetLevel(level int) {
	loggerLevel = level
}

func Info(v ...interface{}) {
	if loggerLevel >= InfoLevel {
		log.Println(v...)
	}
}

func Infof(format string, v ...interface{}) {
	if loggerLevel >= InfoLevel {
		log.Printf(format, v...)
	}
}

func Warning(v ...interface{}) {
	if loggerLevel >= WarningLevel {
		log.Println(v...)
	}
}

func Warningf(format string, v ...interface{}) {
	if loggerLevel >= WarningLevel {
		log.Printf(format, v...)
	}
}

func Debug(v ...interface{}) {
	if loggerLevel >= DebugLevel {
		log.Println(v...)
	}
}

func Debugf(format string, v ...interface{}) {
	if loggerLevel >= DebugLevel {
		log.Printf(format, v...)
	}
}