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

     1  package handler
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  
     9  	"github.com/lingyao2333/mo-zero/core/logx"
    10  	"github.com/lingyao2333/mo-zero/core/stat"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  func init() {
    15  	logx.Disable()
    16  	stat.SetReporter(nil)
    17  }
    18  
    19  func TestBreakerHandlerAccept(t *testing.T) {
    20  	metrics := stat.NewMetrics("unit-test")
    21  	breakerHandler := BreakerHandler(http.MethodGet, "/", metrics)
    22  	handler := breakerHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    23  		w.Header().Set("X-Test", "test")
    24  		_, err := w.Write([]byte("content"))
    25  		assert.Nil(t, err)
    26  	}))
    27  
    28  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    29  	req.Header.Set("X-Test", "test")
    30  	resp := httptest.NewRecorder()
    31  	handler.ServeHTTP(resp, req)
    32  	assert.Equal(t, http.StatusOK, resp.Code)
    33  	assert.Equal(t, "test", resp.Header().Get("X-Test"))
    34  	assert.Equal(t, "content", resp.Body.String())
    35  }
    36  
    37  func TestBreakerHandlerFail(t *testing.T) {
    38  	metrics := stat.NewMetrics("unit-test")
    39  	breakerHandler := BreakerHandler(http.MethodGet, "/", metrics)
    40  	handler := breakerHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    41  		w.WriteHeader(http.StatusBadGateway)
    42  	}))
    43  
    44  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    45  	resp := httptest.NewRecorder()
    46  	handler.ServeHTTP(resp, req)
    47  	assert.Equal(t, http.StatusBadGateway, resp.Code)
    48  }
    49  
    50  func TestBreakerHandler_4XX(t *testing.T) {
    51  	metrics := stat.NewMetrics("unit-test")
    52  	breakerHandler := BreakerHandler(http.MethodGet, "/", metrics)
    53  	handler := breakerHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    54  		w.WriteHeader(http.StatusBadRequest)
    55  	}))
    56  
    57  	for i := 0; i < 1000; i++ {
    58  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    59  		resp := httptest.NewRecorder()
    60  		handler.ServeHTTP(resp, req)
    61  	}
    62  
    63  	const tries = 100
    64  	var pass int
    65  	for i := 0; i < tries; i++ {
    66  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    67  		resp := httptest.NewRecorder()
    68  		handler.ServeHTTP(resp, req)
    69  		if resp.Code == http.StatusBadRequest {
    70  			pass++
    71  		}
    72  	}
    73  
    74  	assert.Equal(t, tries, pass)
    75  }
    76  
    77  func TestBreakerHandlerReject(t *testing.T) {
    78  	metrics := stat.NewMetrics("unit-test")
    79  	breakerHandler := BreakerHandler(http.MethodGet, "/", metrics)
    80  	handler := breakerHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    81  		w.WriteHeader(http.StatusInternalServerError)
    82  	}))
    83  
    84  	for i := 0; i < 1000; i++ {
    85  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    86  		resp := httptest.NewRecorder()
    87  		handler.ServeHTTP(resp, req)
    88  	}
    89  
    90  	var drops int
    91  	for i := 0; i < 100; i++ {
    92  		req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    93  		resp := httptest.NewRecorder()
    94  		handler.ServeHTTP(resp, req)
    95  		if resp.Code == http.StatusServiceUnavailable {
    96  			drops++
    97  		}
    98  	}
    99  
   100  	assert.True(t, drops >= 80, fmt.Sprintf("expected to be greater than 80, but got %d", drops))
   101  }