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

     1  package handler
     2  
     3  import (
     4  	"io"
     5  	"log"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  
    10  	"github.com/lingyao2333/mo-zero/core/load"
    11  	"github.com/lingyao2333/mo-zero/core/stat"
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  func init() {
    16  	log.SetOutput(io.Discard)
    17  }
    18  
    19  func TestSheddingHandlerAccept(t *testing.T) {
    20  	metrics := stat.NewMetrics("unit-test")
    21  	shedder := mockShedder{
    22  		allow: true,
    23  	}
    24  	sheddingHandler := SheddingHandler(shedder, metrics)
    25  	handler := sheddingHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    26  		w.Header().Set("X-Test", "test")
    27  		_, err := w.Write([]byte("content"))
    28  		assert.Nil(t, err)
    29  	}))
    30  
    31  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    32  	req.Header.Set("X-Test", "test")
    33  	resp := httptest.NewRecorder()
    34  	handler.ServeHTTP(resp, req)
    35  	assert.Equal(t, http.StatusOK, resp.Code)
    36  	assert.Equal(t, "test", resp.Header().Get("X-Test"))
    37  	assert.Equal(t, "content", resp.Body.String())
    38  }
    39  
    40  func TestSheddingHandlerFail(t *testing.T) {
    41  	metrics := stat.NewMetrics("unit-test")
    42  	shedder := mockShedder{
    43  		allow: true,
    44  	}
    45  	sheddingHandler := SheddingHandler(shedder, metrics)
    46  	handler := sheddingHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    47  		w.WriteHeader(http.StatusServiceUnavailable)
    48  	}))
    49  
    50  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    51  	resp := httptest.NewRecorder()
    52  	handler.ServeHTTP(resp, req)
    53  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    54  }
    55  
    56  func TestSheddingHandlerReject(t *testing.T) {
    57  	metrics := stat.NewMetrics("unit-test")
    58  	shedder := mockShedder{
    59  		allow: false,
    60  	}
    61  	sheddingHandler := SheddingHandler(shedder, metrics)
    62  	handler := sheddingHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    63  		w.WriteHeader(http.StatusOK)
    64  	}))
    65  
    66  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    67  	resp := httptest.NewRecorder()
    68  	handler.ServeHTTP(resp, req)
    69  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    70  }
    71  
    72  func TestSheddingHandlerNoShedding(t *testing.T) {
    73  	metrics := stat.NewMetrics("unit-test")
    74  	sheddingHandler := SheddingHandler(nil, metrics)
    75  	handler := sheddingHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    76  		w.WriteHeader(http.StatusOK)
    77  	}))
    78  
    79  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    80  	resp := httptest.NewRecorder()
    81  	handler.ServeHTTP(resp, req)
    82  	assert.Equal(t, http.StatusOK, resp.Code)
    83  }
    84  
    85  type mockShedder struct {
    86  	allow bool
    87  }
    88  
    89  func (s mockShedder) Allow() (load.Promise, error) {
    90  	if s.allow {
    91  		return mockPromise{}, nil
    92  	}
    93  
    94  	return nil, load.ErrServiceOverloaded
    95  }
    96  
    97  type mockPromise struct{}
    98  
    99  func (p mockPromise) Pass() {
   100  }
   101  
   102  func (p mockPromise) Fail() {
   103  }