sheddinghandler_test.go 2.6 KB

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