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

     1  package handler
     2  
     3  import (
     4  	"bytes"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  func TestMaxBytesHandler(t *testing.T) {
    13  	maxb := MaxBytesHandler(10)
    14  	handler := maxb(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
    15  
    16  	req := httptest.NewRequest(http.MethodPost, "http://localhost",
    17  		bytes.NewBufferString("123456789012345"))
    18  	resp := httptest.NewRecorder()
    19  	handler.ServeHTTP(resp, req)
    20  	assert.Equal(t, http.StatusRequestEntityTooLarge, resp.Code)
    21  
    22  	req = httptest.NewRequest(http.MethodPost, "http://localhost", bytes.NewBufferString("12345"))
    23  	resp = httptest.NewRecorder()
    24  	handler.ServeHTTP(resp, req)
    25  	assert.Equal(t, http.StatusOK, resp.Code)
    26  }
    27  
    28  func TestMaxBytesHandlerNoLimit(t *testing.T) {
    29  	maxb := MaxBytesHandler(-1)
    30  	handler := maxb(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
    31  
    32  	req := httptest.NewRequest(http.MethodPost, "http://localhost",
    33  		bytes.NewBufferString("123456789012345"))
    34  	resp := httptest.NewRecorder()
    35  	handler.ServeHTTP(resp, req)
    36  	assert.Equal(t, http.StatusOK, resp.Code)
    37  }