From 93124329acffcd3016f6131753117f2a3ed598de Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 8 May 2023 23:49:13 +0800
Subject: [PATCH] chore: add more tests (#3229)

---
 core/rescue/recover.go      |  4 ++--
 core/rescue/recover_test.go | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/core/rescue/recover.go b/core/rescue/recover.go
index 5d2922bf..b9361b7c 100644
--- a/core/rescue/recover.go
+++ b/core/rescue/recover.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"runtime/debug"
 
-	"github.com/zeromicro/go-zero/core/logc"
 	"github.com/zeromicro/go-zero/core/logx"
 )
 
@@ -22,12 +21,13 @@ func Recover(cleanups ...func()) {
 	}
 }
 
+// RecoverCtx is used with defer to do cleanup on panics.
 func RecoverCtx(ctx context.Context, cleanups ...func()) {
 	for _, cleanup := range cleanups {
 		cleanup()
 	}
 
 	if p := recover(); p != nil {
-		logc.Errorf(ctx, "%+v\n\n%s", p, debug.Stack())
+		logx.WithContext(ctx).Errorf("%+v\n%s", p, debug.Stack())
 	}
 }
diff --git a/core/rescue/recover_test.go b/core/rescue/recover_test.go
index fa4cd33c..07816054 100644
--- a/core/rescue/recover_test.go
+++ b/core/rescue/recover_test.go
@@ -1,6 +1,7 @@
 package rescue
 
 import (
+	"context"
 	"sync/atomic"
 	"testing"
 
@@ -25,3 +26,17 @@ func TestRescue(t *testing.T) {
 	})
 	assert.Equal(t, int32(5), atomic.LoadInt32(&count))
 }
+
+func TestRescueCtx(t *testing.T) {
+	var count int32
+	assert.NotPanics(t, func() {
+		defer RecoverCtx(context.Background(), func() {
+			atomic.AddInt32(&count, 2)
+		}, func() {
+			atomic.AddInt32(&count, 3)
+		})
+
+		panic("hello")
+	})
+	assert.Equal(t, int32(5), atomic.LoadInt32(&count))
+}