github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/loghandler_test.go (about)

     1  package handler
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"log"
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/lingyao2333/mo-zero/rest/internal"
    13  	"github.com/stretchr/testify/assert"
    14  )
    15  
    16  func init() {
    17  	log.SetOutput(io.Discard)
    18  }
    19  
    20  func TestLogHandler(t *testing.T) {
    21  	handlers := []func(handler http.Handler) http.Handler{
    22  		LogHandler,
    23  		DetailedLogHandler,
    24  	}
    25  
    26  	for _, logHandler := range handlers {
    27  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    28  		handler := logHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    29  			r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything")
    30  			w.Header().Set("X-Test", "test")
    31  			w.WriteHeader(http.StatusServiceUnavailable)
    32  			_, err := w.Write([]byte("content"))
    33  			assert.Nil(t, err)
    34  
    35  			flusher, ok := w.(http.Flusher)
    36  			assert.True(t, ok)
    37  			flusher.Flush()
    38  		}))
    39  
    40  		resp := httptest.NewRecorder()
    41  		handler.ServeHTTP(resp, req)
    42  		assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    43  		assert.Equal(t, "test", resp.Header().Get("X-Test"))
    44  		assert.Equal(t, "content", resp.Body.String())
    45  	}
    46  }
    47  
    48  func TestLogHandlerVeryLong(t *testing.T) {
    49  	var buf bytes.Buffer
    50  	for i := 0; i < limitBodyBytes<<1; i++ {
    51  		buf.WriteByte('a')
    52  	}
    53  
    54  	req := httptest.NewRequest(http.MethodPost, "http://localhost", &buf)
    55  	handler := LogHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    56  		r.Context().Value(internal.LogContext).(*internal.LogCollector).Append("anything")
    57  		io.Copy(io.Discard, r.Body)
    58  		w.Header().Set("X-Test", "test")
    59  		w.WriteHeader(http.StatusServiceUnavailable)
    60  		_, err := w.Write([]byte("content"))
    61  		assert.Nil(t, err)
    62  
    63  		flusher, ok := w.(http.Flusher)
    64  		assert.True(t, ok)
    65  		flusher.Flush()
    66  	}))
    67  
    68  	resp := httptest.NewRecorder()
    69  	handler.ServeHTTP(resp, req)
    70  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    71  	assert.Equal(t, "test", resp.Header().Get("X-Test"))
    72  	assert.Equal(t, "content", resp.Body.String())
    73  }
    74  
    75  func TestLogHandlerSlow(t *testing.T) {
    76  	handlers := []func(handler http.Handler) http.Handler{
    77  		LogHandler,
    78  		DetailedLogHandler,
    79  	}
    80  
    81  	for _, logHandler := range handlers {
    82  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    83  		handler := logHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    84  			time.Sleep(defaultSlowThreshold + time.Millisecond*50)
    85  		}))
    86  
    87  		resp := httptest.NewRecorder()
    88  		handler.ServeHTTP(resp, req)
    89  		assert.Equal(t, http.StatusOK, resp.Code)
    90  	}
    91  }
    92  
    93  func TestLogHandler_Hijack(t *testing.T) {
    94  	resp := httptest.NewRecorder()
    95  	writer := &loggedResponseWriter{
    96  		w: resp,
    97  	}
    98  	assert.NotPanics(t, func() {
    99  		writer.Hijack()
   100  	})
   101  
   102  	writer = &loggedResponseWriter{
   103  		w: mockedHijackable{resp},
   104  	}
   105  	assert.NotPanics(t, func() {
   106  		writer.Hijack()
   107  	})
   108  }
   109  
   110  func TestDetailedLogHandler_Hijack(t *testing.T) {
   111  	resp := httptest.NewRecorder()
   112  	writer := &detailLoggedResponseWriter{
   113  		writer: &loggedResponseWriter{
   114  			w: resp,
   115  		},
   116  	}
   117  	assert.NotPanics(t, func() {
   118  		writer.Hijack()
   119  	})
   120  
   121  	writer = &detailLoggedResponseWriter{
   122  		writer: &loggedResponseWriter{
   123  			w: mockedHijackable{resp},
   124  		},
   125  	}
   126  	assert.NotPanics(t, func() {
   127  		writer.Hijack()
   128  	})
   129  }
   130  
   131  func TestSetSlowThreshold(t *testing.T) {
   132  	assert.Equal(t, defaultSlowThreshold, slowThreshold.Load())
   133  	SetSlowThreshold(time.Second)
   134  	assert.Equal(t, time.Second, slowThreshold.Load())
   135  }
   136  
   137  func TestWrapMethodWithColor(t *testing.T) {
   138  	// no tty
   139  	assert.Equal(t, http.MethodGet, wrapMethod(http.MethodGet))
   140  	assert.Equal(t, http.MethodPost, wrapMethod(http.MethodPost))
   141  	assert.Equal(t, http.MethodPut, wrapMethod(http.MethodPut))
   142  	assert.Equal(t, http.MethodDelete, wrapMethod(http.MethodDelete))
   143  	assert.Equal(t, http.MethodPatch, wrapMethod(http.MethodPatch))
   144  	assert.Equal(t, http.MethodHead, wrapMethod(http.MethodHead))
   145  	assert.Equal(t, http.MethodOptions, wrapMethod(http.MethodOptions))
   146  	assert.Equal(t, http.MethodConnect, wrapMethod(http.MethodConnect))
   147  	assert.Equal(t, http.MethodTrace, wrapMethod(http.MethodTrace))
   148  }
   149  
   150  func TestWrapStatusCodeWithColor(t *testing.T) {
   151  	// no tty
   152  	assert.Equal(t, "200", wrapStatusCode(http.StatusOK))
   153  	assert.Equal(t, "302", wrapStatusCode(http.StatusFound))
   154  	assert.Equal(t, "404", wrapStatusCode(http.StatusNotFound))
   155  	assert.Equal(t, "500", wrapStatusCode(http.StatusInternalServerError))
   156  	assert.Equal(t, "503", wrapStatusCode(http.StatusServiceUnavailable))
   157  }
   158  
   159  func BenchmarkLogHandler(b *testing.B) {
   160  	b.ReportAllocs()
   161  
   162  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
   163  	handler := LogHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   164  		w.WriteHeader(http.StatusOK)
   165  	}))
   166  
   167  	for i := 0; i < b.N; i++ {
   168  		resp := httptest.NewRecorder()
   169  		handler.ServeHTTP(resp, req)
   170  	}
   171  }