chizap/logger_test.go

package chizap_test

import (
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.dwysokinski.me/Kichiyaki/chizap"
	"github.com/go-chi/chi/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"
)

func TestLogger(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name                     string
		req                      *http.Request
		excluded                 bool
		expectedIP               string
		expectedLevel            zapcore.Level
		expectedRoutePattern     string
		expectedAdditionalFields []zap.Field
	}{
		{
			name:                 "/info?test=true",
			req:                  httptest.NewRequest(http.MethodGet, "/info?test=true", nil),
			expectedIP:           "192.0.2.1",
			expectedLevel:        zap.InfoLevel,
			expectedRoutePattern: "/info",
		},
		{
			name:                 "/warn?test=true",
			req:                  httptest.NewRequest(http.MethodGet, "/warn?test=true", nil),
			expectedIP:           "192.0.2.1",
			expectedLevel:        zap.WarnLevel,
			expectedRoutePattern: "/warn",
		},
		{
			name:                 "/error?test=true",
			req:                  httptest.NewRequest(http.MethodGet, "/error?test=true", nil),
			expectedIP:           "192.0.2.1",
			expectedLevel:        zap.ErrorLevel,
			expectedRoutePattern: "/error",
		},
		{
			name:     "/excluded?test=true",
			req:      httptest.NewRequest(http.MethodGet, "/excluded?test=true", nil),
			excluded: true,
		},
		{
			name:                 "/delete/123",
			req:                  httptest.NewRequest(http.MethodDelete, "/delete/123", nil),
			expectedIP:           "192.0.2.1",
			expectedLevel:        zap.InfoLevel,
			expectedRoutePattern: "/delete/{id}",
			expectedAdditionalFields: []zap.Field{
				zap.String("id", "123"),
			},
		},
		{
			name: "/x-forwarded-for",
			req: func() *http.Request {
				req := httptest.NewRequest(http.MethodGet, "/x-forwarded-for", nil)
				req.Header.Set("X-Forwarded-For", "94.222.111.115")
				return req
			}(),
			expectedIP:           "94.222.111.115",
			expectedLevel:        zap.InfoLevel,
			expectedRoutePattern: "/x-forwarded-for",
		},
	}

	for _, tt := range tests {
		tt := tt

		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			logger, obs := newLogger()

			rr := httptest.NewRecorder()
			newRouter(logger).ServeHTTP(rr, tt.req)

			entries := obs.TakeAll()

			if tt.excluded {
				require.Len(t, entries, 0)
				return
			}

			require.Len(t, entries, 1)
			entry := entries[0]
			assert.Equal(t, tt.req.URL.Path, entry.Message)
			assert.Equal(t, tt.expectedLevel, entry.Level)
			require.Len(t, entry.Context, 11+len(tt.expectedAdditionalFields))
			assert.Contains(t, entry.Context, zap.String("ip", tt.expectedIP))
			assert.Contains(t, entry.Context, zap.Int("statusCode", rr.Code))
			assert.Contains(t, entry.Context, zap.String("method", tt.req.Method))
			assert.Contains(t, entry.Context, zap.String("path", tt.req.URL.Path))
			assert.Contains(t, entry.Context, zap.String("query", tt.req.URL.RawQuery))
			assert.Contains(t, entry.Context, zap.String("proto", tt.req.Proto))
			assert.Contains(t, entry.Context, zap.String("referer", tt.req.Referer()))
			assert.Contains(t, entry.Context, zap.String("userAgent", tt.req.UserAgent()))
			assert.Contains(t, entry.Context, zap.String("routePattern", tt.expectedRoutePattern))
			for _, f := range tt.expectedAdditionalFields {
				assert.Contains(t, entry.Context, f)
			}
		})
	}
}

func newLogger() (*zap.Logger, *observer.ObservedLogs) {
	core, obs := observer.New(zap.NewAtomicLevelAt(zap.InfoLevel))
	logger := zap.New(core)
	return logger, obs
}

func newRouter(logger *zap.Logger) *chi.Mux {
	router := chi.NewRouter()

	router.Route("/", func(r chi.Router) {
		r.Use(chizap.Logger(
			logger,
			chizap.WithFilter(func(r *http.Request) bool {
				return r.URL.Path != "/excluded"
			}),
			chizap.WithFilter(func(r *http.Request) bool {
				return r.URL.Path != "/excluded2"
			}),
			chizap.WithAdditionalFieldExtractor(func(r *http.Request) []zap.Field {
				if !strings.HasPrefix(r.URL.Path, "/delete") {
					return nil
				}

				return []zap.Field{
					zap.String("id", chi.URLParam(r, "id")),
				}
			}),
			chizap.WithIPFn(func(r *http.Request) string {
				if r.URL.Path != "/x-forwarded-for" {
					ip, _, _ := net.SplitHostPort(r.RemoteAddr)
					return ip
				}
				return r.Header.Get("X-Forwarded-For")
			}),
		))
		r.Get("/info", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})
		r.Get("/warn", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusBadRequest)
		})
		r.Get("/error", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusInternalServerError)
		})
		r.Get("/excluded", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})
		r.Get("/x-forwarded-for", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})
		r.Delete("/delete/{id}", func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusOK)
		})
	})

	return router
}