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)

		if loggerLevel >= InfoLevel {
			logrus.WithFields(logrus.Fields{
				"URI":      req.RequestURI,
				"Method":   req.Method,
				"Status":   responseData.status,
				"Duration": duration,
				"Size":     responseData.size,
			}).Info("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...)
	}
}