test: adjust feature flag version guard test

util/featureflag/featureflag_test.go

@@ -8,9 +8,9 @@ import (
 	"testing"
 
 	"github.com/gin-gonic/gin"
-	"github.com/joomcode/errorx"
-	"github.com/pingcap/tidb-dashboard/util/rest"
 	"github.com/stretchr/testify/require"
+
+	"github.com/pingcap/tidb-dashboard/util/rest"
 )
 
 func Test_Name(t *testing.T) {
@@ -54,39 +54,25 @@ func Test_VersionGuard(t *testing.T) {
 	e := gin.Default()
 	e.Use(f1.VersionGuard())
 	e.GET("/ping", func(c *gin.Context) {
-		c.String(200, "pong")
+		c.String(http.StatusOK, "pong")
 	})
 	w := httptest.NewRecorder()
 	req, _ := http.NewRequest("GET", "/ping", nil)
 	e.ServeHTTP(w, req)
 
-	r.Equal(200, w.Code)
+	r.Equal(http.StatusOK, w.Code)
 	r.Equal("pong", w.Body.String())
 
-	// abort with other middlewares
-	handled := false
+	// abort with rest.ErrorHandlerFn
 	e2 := gin.Default()
-	e2.Use(func(c *gin.Context) {
-		c.Next()
-
-		handled = true
-
-		// check error type
-		currentErr := c.Errors.Last().Err
-		codeProperty, _ := rest.HTTPCodeProperty(http.StatusForbidden)
-		code, ok := currentErr.(*errorx.Error).Property(codeProperty)
-		r.True(ok)
-
-		r.Equal(true, errorx.IsOfType(currentErr, ErrFeatureUnsupported))
-		r.Equal(http.StatusForbidden, code)
-	})
+	e2.Use(rest.ErrorHandlerFn())
 	e2.Use(f1.VersionGuard(), f2.VersionGuard())
 	e2.GET("/ping", func(c *gin.Context) {
-		c.String(200, "pong")
+		c.String(http.StatusOK, "pong")
 	})
 	w2 := httptest.NewRecorder()
 	req2, _ := http.NewRequest("GET", "/ping", nil)
 	e2.ServeHTTP(w2, req2)
 
-	r.Equal(true, handled)
+	r.Equal(http.StatusForbidden, w2.Code)
 }