close #4483: fix failpoint

Signed-off-by: Cabinfever_B <cabinfeveroier@gmail.com>
tikv · Dec 22, 2021 · 9fef69a · 9fef69a
1 parent dd14b0b
commit 9fef69a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/pkg/apiutil/serverapi/middleware.go b/pkg/apiutil/serverapi/middleware.go
@@ -91,11 +91,12 @@ func NewSelfProtector(s *server.Server) negroni.Handler {
 
 func (protector *selfProtector) ServeHTTP(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
 	handler := protector.s.GetSelfProtectionHandler()
-	if handler == nil || handler.HandleHTTPSelfProtection(r) {
-		failpoint.Inject("addSelfProtectionHTTPHeader", func() {
-			w.Header().Add("self-protection", "ok")
-		})
 
+	failpoint.Inject("addSelfProtectionHTTPHeader", func() {
+		w.Header().Add("self-protection", "ok")
+	})
+
+	if handler == nil || handler.HandleHTTPSelfProtection(r) {
 		next(w, r)
 	} else {
 		// current plan will only deny request when over the speed limit

diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go
@@ -135,9 +135,9 @@ func (s *testSelfProtectorSuite) TearDownSuite(c *C) {
 }
 
 func (s *testSelfProtectorSuite) TestSelfProtect(c *C) {
+	c.Assert(failpoint.Enable("github.com/tikv/pd/pkg/apiutil/serverapi/addSelfProtectionHTTPHeader", "return(true)"), IsNil)
 	leader := s.cluster.GetServer(s.cluster.GetLeader())
 	header := mustRequestSuccess(c, leader.GetServer())
-	c.Assert(failpoint.Enable("github.com/tikv/pd/pkg/apiutil/serverapi/addSelfProtectionHTTPHeader", `return(true)`), IsNil)
 	c.Assert(header.Get("self-protection"), Equals, "ok")
 	c.Assert(failpoint.Disable("github.com/tikv/pd/pkg/apiutil/serverapi/addSelfProtectionHTTPHeader"), IsNil)
 }