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

     1  package handler
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"log"
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  func init() {
    16  	log.SetOutput(io.Discard)
    17  }
    18  
    19  func TestTimeout(t *testing.T) {
    20  	timeoutHandler := TimeoutHandler(time.Millisecond)
    21  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    22  		time.Sleep(time.Minute)
    23  	}))
    24  
    25  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    26  	resp := httptest.NewRecorder()
    27  	handler.ServeHTTP(resp, req)
    28  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    29  }
    30  
    31  func TestWithinTimeout(t *testing.T) {
    32  	timeoutHandler := TimeoutHandler(time.Second)
    33  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    34  		time.Sleep(time.Millisecond)
    35  	}))
    36  
    37  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    38  	resp := httptest.NewRecorder()
    39  	handler.ServeHTTP(resp, req)
    40  	assert.Equal(t, http.StatusOK, resp.Code)
    41  }
    42  
    43  func TestWithTimeoutTimedout(t *testing.T) {
    44  	timeoutHandler := TimeoutHandler(time.Millisecond)
    45  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    46  		time.Sleep(time.Millisecond * 10)
    47  		w.Write([]byte(`foo`))
    48  		w.WriteHeader(http.StatusOK)
    49  	}))
    50  
    51  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    52  	resp := httptest.NewRecorder()
    53  	handler.ServeHTTP(resp, req)
    54  	assert.Equal(t, http.StatusServiceUnavailable, resp.Code)
    55  }
    56  
    57  func TestWithoutTimeout(t *testing.T) {
    58  	timeoutHandler := TimeoutHandler(0)
    59  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    60  		time.Sleep(100 * time.Millisecond)
    61  	}))
    62  
    63  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    64  	resp := httptest.NewRecorder()
    65  	handler.ServeHTTP(resp, req)
    66  	assert.Equal(t, http.StatusOK, resp.Code)
    67  }
    68  
    69  func TestTimeoutPanic(t *testing.T) {
    70  	timeoutHandler := TimeoutHandler(time.Minute)
    71  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    72  		panic("foo")
    73  	}))
    74  
    75  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    76  	resp := httptest.NewRecorder()
    77  	assert.Panics(t, func() {
    78  		handler.ServeHTTP(resp, req)
    79  	})
    80  }
    81  
    82  func TestTimeoutWebsocket(t *testing.T) {
    83  	timeoutHandler := TimeoutHandler(time.Millisecond)
    84  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    85  		time.Sleep(time.Millisecond * 10)
    86  	}))
    87  
    88  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
    89  	req.Header.Set(headerUpgrade, valueWebsocket)
    90  	resp := httptest.NewRecorder()
    91  	handler.ServeHTTP(resp, req)
    92  	assert.Equal(t, http.StatusOK, resp.Code)
    93  }
    94  
    95  func TestTimeoutWroteHeaderTwice(t *testing.T) {
    96  	timeoutHandler := TimeoutHandler(time.Minute)
    97  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    98  		w.Write([]byte(`hello`))
    99  		w.Header().Set("foo", "bar")
   100  		w.WriteHeader(http.StatusOK)
   101  	}))
   102  
   103  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
   104  	resp := httptest.NewRecorder()
   105  	handler.ServeHTTP(resp, req)
   106  	assert.Equal(t, http.StatusOK, resp.Code)
   107  }
   108  
   109  func TestTimeoutWriteBadCode(t *testing.T) {
   110  	timeoutHandler := TimeoutHandler(time.Minute)
   111  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   112  		w.WriteHeader(1000)
   113  	}))
   114  
   115  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
   116  	resp := httptest.NewRecorder()
   117  	assert.Panics(t, func() {
   118  		handler.ServeHTTP(resp, req)
   119  	})
   120  }
   121  
   122  func TestTimeoutClientClosed(t *testing.T) {
   123  	timeoutHandler := TimeoutHandler(time.Minute)
   124  	handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   125  		w.WriteHeader(http.StatusServiceUnavailable)
   126  	}))
   127  
   128  	req := httptest.NewRequest(http.MethodGet, "http://localhost", http.NoBody)
   129  	ctx, cancel := context.WithCancel(context.Background())
   130  	req = req.WithContext(ctx)
   131  	cancel()
   132  	resp := httptest.NewRecorder()
   133  	handler.ServeHTTP(resp, req)
   134  	assert.Equal(t, statusClientClosedRequest, resp.Code)
   135  }
   136  
   137  func TestTimeoutPusher(t *testing.T) {
   138  	handler := &timeoutWriter{
   139  		w: mockedPusher{},
   140  	}
   141  
   142  	assert.Panics(t, func() {
   143  		handler.Push("any", nil)
   144  	})
   145  
   146  	handler = &timeoutWriter{
   147  		w: httptest.NewRecorder(),
   148  	}
   149  	assert.Equal(t, http.ErrNotSupported, handler.Push("any", nil))
   150  }
   151  
   152  type mockedPusher struct{}
   153  
   154  func (m mockedPusher) Header() http.Header {
   155  	panic("implement me")
   156  }
   157  
   158  func (m mockedPusher) Write(bytes []byte) (int, error) {
   159  	panic("implement me")
   160  }
   161  
   162  func (m mockedPusher) WriteHeader(statusCode int) {
   163  	panic("implement me")
   164  }
   165  
   166  func (m mockedPusher) Push(target string, opts *http.PushOptions) error {
   167  	panic("implement me")
   168  }