Browse Source

chore: add more tests (#3229)

Kevin Wan 2 years ago
parent
commit
93124329ac
2 changed files with 17 additions and 2 deletions
  1. 2 2
      core/rescue/recover.go
  2. 15 0
      core/rescue/recover_test.go

+ 2 - 2
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())
 	}
 }

+ 15 - 0
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))
+}