2024-09-24 21:07:42 +00:00
|
|
|
package middleware
|
|
|
|
|
|
|
|
import (
|
2024-09-25 11:52:12 +00:00
|
|
|
"bytes"
|
|
|
|
"io"
|
2024-09-24 21:07:42 +00:00
|
|
|
"net/http"
|
|
|
|
"net/http/httptest"
|
|
|
|
"testing"
|
|
|
|
"time"
|
|
|
|
|
|
|
|
"git.vinchent.xyz/vinchent/go-web/framework"
|
|
|
|
)
|
|
|
|
|
|
|
|
func TestTimeout(t *testing.T) {
|
|
|
|
t.Run("Test timeout handler", func(t *testing.T) {
|
|
|
|
timeoutHandler := Timeout(1 * time.Millisecond)
|
|
|
|
|
2024-09-25 11:52:12 +00:00
|
|
|
longHandler := func(c *framework.Context) error {
|
|
|
|
time.Sleep(2 * time.Millisecond)
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2024-09-25 12:09:29 +00:00
|
|
|
res := prepareMiddlewareTest(t, timeoutHandler, longHandler)
|
2024-09-25 11:52:12 +00:00
|
|
|
|
2024-09-25 12:09:29 +00:00
|
|
|
assertCode(t, res.StatusCode, http.StatusRequestTimeout)
|
|
|
|
assertBody(t, res.Body, "time out")
|
|
|
|
})
|
|
|
|
t.Run("Test no timeout", func(t *testing.T) {
|
|
|
|
timeoutHandler := Timeout(2 * time.Millisecond)
|
|
|
|
|
|
|
|
quickHandler := func(c *framework.Context) error {
|
|
|
|
// time.Sleep(1 * time.Millisecond)
|
|
|
|
c.WriteJSON(http.StatusOK, "ok")
|
|
|
|
return nil
|
2024-09-25 11:52:12 +00:00
|
|
|
}
|
2024-09-25 12:09:29 +00:00
|
|
|
|
|
|
|
res := prepareMiddlewareTest(t, timeoutHandler, quickHandler)
|
|
|
|
|
|
|
|
assertCode(t, res.StatusCode, http.StatusOK)
|
2024-09-24 21:07:42 +00:00
|
|
|
})
|
|
|
|
}
|
2024-09-25 11:52:12 +00:00
|
|
|
|
|
|
|
func TestRecover(t *testing.T) {
|
2024-09-25 12:09:29 +00:00
|
|
|
t.Run("Test panic", func(t *testing.T) {
|
|
|
|
recoverer := Recovery()
|
2024-09-25 11:52:12 +00:00
|
|
|
|
2024-09-25 12:09:29 +00:00
|
|
|
panicHandler := func(c *framework.Context) error {
|
|
|
|
panic("panic")
|
|
|
|
}
|
|
|
|
|
|
|
|
res := prepareMiddlewareTest(t, recoverer, panicHandler)
|
|
|
|
|
|
|
|
assertCode(t, res.StatusCode, http.StatusInternalServerError)
|
|
|
|
})
|
|
|
|
t.Run("Test no panic", func(t *testing.T) {
|
|
|
|
recoverer := Recovery()
|
|
|
|
|
|
|
|
normalHandler := func(c *framework.Context) error {
|
|
|
|
c.WriteJSON(http.StatusOK, "ok")
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
res := prepareMiddlewareTest(t, recoverer, normalHandler)
|
|
|
|
|
|
|
|
assertCode(t, res.StatusCode, http.StatusOK)
|
|
|
|
})
|
|
|
|
}
|
|
|
|
|
|
|
|
func prepareMiddlewareTest(
|
|
|
|
t testing.TB,
|
|
|
|
mid framework.ControllerHandler,
|
|
|
|
in framework.ControllerHandler,
|
|
|
|
) *http.Response {
|
|
|
|
t.Helper()
|
2024-09-25 11:52:12 +00:00
|
|
|
request := httptest.NewRequest(http.MethodGet, "/", nil)
|
|
|
|
response := httptest.NewRecorder()
|
|
|
|
c := framework.NewContext(response, request)
|
|
|
|
|
2024-09-25 12:09:29 +00:00
|
|
|
c.SetHandlers([]framework.ControllerHandler{in})
|
2024-09-25 11:52:12 +00:00
|
|
|
|
2024-09-25 12:09:29 +00:00
|
|
|
err := mid(c)
|
2024-09-25 11:52:12 +00:00
|
|
|
if err != nil {
|
|
|
|
t.Fatal(err)
|
|
|
|
}
|
2024-09-25 12:09:29 +00:00
|
|
|
|
2024-09-25 11:52:12 +00:00
|
|
|
res := response.Result()
|
2024-09-25 12:09:29 +00:00
|
|
|
return res
|
|
|
|
}
|
|
|
|
|
|
|
|
func assertCode(t testing.TB, got int, want int) {
|
|
|
|
t.Helper()
|
|
|
|
if got != want {
|
|
|
|
t.Errorf("status code got %d, want %d", got, want)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func assertBody(t testing.TB, got io.Reader, want string) {
|
|
|
|
t.Helper()
|
|
|
|
buf, _ := io.ReadAll(got)
|
|
|
|
if cmp := bytes.Compare(buf, []byte(want)); cmp != 0 {
|
|
|
|
t.Errorf("got %q, want %q", string(buf), want)
|
2024-09-25 11:52:12 +00:00
|
|
|
}
|
|
|
|
}
|