diff --git a/recovery.go b/recovery.go index 2f4caae0..d02b829b 100644 --- a/recovery.go +++ b/recovery.go @@ -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) } diff --git a/recovery_test.go b/recovery_test.go index ed8f6489..6cc2a47a 100644 --- a/recovery_test.go +++ b/recovery_test.go @@ -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