support RecoveryWithWriter(io.Writer, ...RecoveryFunc)

This commit is contained in:
Johnny Dallas 2020-06-12 10:22:53 +00:00
parent ad35b14568
commit 5782c13d25
2 changed files with 43 additions and 5 deletions

View File

@ -36,11 +36,14 @@ func Recovery() HandlerFunc {
//CustomRecovery returns a middleware that recovers from any panics and calls the provided handle func to handle it.
func CustomRecovery(handle RecoveryFunc) HandlerFunc {
return CustomRecoveryWithWriter(DefaultErrorWriter, handle)
return RecoveryWithWriter(DefaultErrorWriter, handle)
}
// RecoveryWithWriter returns a middleware for a given writer that recovers from any panics and writes a 500 if there was one.
func RecoveryWithWriter(out io.Writer) HandlerFunc {
func RecoveryWithWriter(out io.Writer, recovery ...RecoveryFunc) HandlerFunc {
if len(recovery) > 0 {
return CustomRecoveryWithWriter(out, recovery[0])
}
return CustomRecoveryWithWriter(out, defaultHandleRecovery)
}

View File

@ -62,7 +62,7 @@ func TestPanicInHandler(t *testing.T) {
assert.Equal(t, http.StatusInternalServerError, w.Code)
assert.Contains(t, buffer.String(), "panic recovered")
assert.Contains(t, buffer.String(), "Oupps, Houston, we have a problem")
assert.Contains(t, buffer.String(), "TestPanicInHandler")
assert.Contains(t, buffer.String(), t.Name())
assert.NotContains(t, buffer.String(), "GET /recovery")
// Debug mode prints the request
@ -163,7 +163,7 @@ func TestCustomRecoveryWithWriter(t *testing.T) {
assert.Equal(t, http.StatusBadRequest, w.Code)
assert.Contains(t, buffer.String(), "panic recovered")
assert.Contains(t, buffer.String(), "Oupps, Houston, we have a problem")
assert.Contains(t, buffer.String(), "TestCustomRecoveryWithWriter")
assert.Contains(t, buffer.String(), t.Name())
assert.NotContains(t, buffer.String(), "GET /recovery")
// Debug mode prints the request
@ -198,7 +198,42 @@ func TestCustomRecovery(t *testing.T) {
assert.Equal(t, http.StatusBadRequest, w.Code)
assert.Contains(t, buffer.String(), "panic recovered")
assert.Contains(t, buffer.String(), "Oupps, Houston, we have a problem")
assert.Contains(t, buffer.String(), "TestCustomRecovery")
assert.Contains(t, buffer.String(), t.Name())
assert.NotContains(t, buffer.String(), "GET /recovery")
// Debug mode prints the request
SetMode(DebugMode)
// RUN
w = performRequest(router, "GET", "/recovery")
// TEST
assert.Equal(t, http.StatusBadRequest, w.Code)
assert.Contains(t, buffer.String(), "GET /recovery")
assert.Equal(t, strings.Repeat("Oupps, Houston, we have a problem", 2), errBuffer.String())
SetMode(TestMode)
}
func TestRecoveryWithWriterWithCustomRecovery(t *testing.T) {
errBuffer := new(bytes.Buffer)
buffer := new(bytes.Buffer)
router := New()
DefaultErrorWriter = buffer
handleRecovery := func(c *Context, err interface{}) {
errBuffer.WriteString(err.(string))
c.AbortWithStatus(http.StatusBadRequest)
}
router.Use(RecoveryWithWriter(DefaultErrorWriter, handleRecovery))
router.GET("/recovery", func(_ *Context) {
panic("Oupps, Houston, we have a problem")
})
// RUN
w := performRequest(router, "GET", "/recovery")
// TEST
assert.Equal(t, http.StatusBadRequest, w.Code)
assert.Contains(t, buffer.String(), "panic recovered")
assert.Contains(t, buffer.String(), "Oupps, Houston, we have a problem")
assert.Contains(t, buffer.String(), t.Name())
assert.NotContains(t, buffer.String(), "GET /recovery")
// Debug mode prints the request