package handler import ( "io/ioutil" "log" "net/http" "net/http/httptest" "testing" "time" "github.com/stretchr/testify/assert" ) func init() { log.SetOutput(ioutil.Discard) } func TestTimeout(t *testing.T) { timeoutHandler := TimeoutHandler(time.Millisecond) handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(time.Minute) })) req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) resp := httptest.NewRecorder() handler.ServeHTTP(resp, req) assert.Equal(t, http.StatusServiceUnavailable, resp.Code) } func TestWithinTimeout(t *testing.T) { timeoutHandler := TimeoutHandler(time.Second) handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(time.Millisecond) })) req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) resp := httptest.NewRecorder() handler.ServeHTTP(resp, req) assert.Equal(t, http.StatusOK, resp.Code) } func TestWithoutTimeout(t *testing.T) { timeoutHandler := TimeoutHandler(0) handler := timeoutHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(100 * time.Millisecond) })) req := httptest.NewRequest(http.MethodGet, "http://localhost", nil) resp := httptest.NewRecorder() handler.ServeHTTP(resp, req) assert.Equal(t, http.StatusOK, resp.Code) }