Merge branch 'master' into refactor-token

tikv · May 15, 2024 · a2caba3 · a2caba3
2 parents c6673ab + ba7e63a
commit a2caba3
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 19 deletions.
diff --git a/tests/server/api/api_test.go b/tests/server/api/api_test.go
@@ -156,15 +156,15 @@ func (suite *middlewareTestSuite) TestRequestInfoMiddleware() {
 	labels := make(map[string]any)
 	labels["testkey"] = "testvalue"
 	data, _ = json.Marshal(labels)
-	resp, err = dialClient.Post(leader.GetAddr()+"/pd/api/v1/debug/pprof/profile?force=true", "application/json", bytes.NewBuffer(data))
+	resp, err = dialClient.Post(leader.GetAddr()+"/pd/api/v1/debug/pprof/profile?seconds=1", "application/json", bytes.NewBuffer(data))
 	re.NoError(err)
 	_, err = io.ReadAll(resp.Body)
 	resp.Body.Close()
 	re.NoError(err)
 	re.Equal(http.StatusOK, resp.StatusCode)
 
 	re.Equal("Profile", resp.Header.Get("service-label"))
-	re.Equal("{\"force\":[\"true\"]}", resp.Header.Get("url-param"))
+	re.Equal("{\"seconds\":[\"1\"]}", resp.Header.Get("url-param"))
 	re.Equal("{\"testkey\":\"testvalue\"}", resp.Header.Get("body-param"))
 	re.Equal("HTTP/1.1/POST:/pd/api/v1/debug/pprof/profile", resp.Header.Get("method"))
 	re.Equal("anonymous", resp.Header.Get("caller-id"))
@@ -182,7 +182,7 @@ func (suite *middlewareTestSuite) TestRequestInfoMiddleware() {
 	re.False(leader.GetServer().GetServiceMiddlewarePersistOptions().IsAuditEnabled())
 
 	header := mustRequestSuccess(re, leader.GetServer())
-	re.Equal("", header.Get("service-label"))
+	re.Equal("GetVersion", header.Get("service-label"))
 
 	re.NoError(failpoint.Disable("github.com/tikv/pd/server/api/addRequestInfoMiddleware"))
 }
@@ -373,6 +373,18 @@ func (suite *middlewareTestSuite) TestRateLimitMiddleware() {
 		re.NoError(err)
 		re.Equal(http.StatusOK, resp.StatusCode)
 	}
+
+	// reset rate limit
+	input = map[string]any{
+		"enable-rate-limit": "true",
+	}
+	data, err = json.Marshal(input)
+	re.NoError(err)
+	req, _ = http.NewRequest(http.MethodPost, leader.GetAddr()+"/pd/api/v1/service-middleware/config", bytes.NewBuffer(data))
+	resp, err = dialClient.Do(req)
+	re.NoError(err)
+	resp.Body.Close()
+	re.True(leader.GetServer().GetServiceMiddlewarePersistOptions().IsRateLimitEnabled())
 }
 
 func (suite *middlewareTestSuite) TestSwaggerUrl() {

diff --git a/tests/server/api/operator_test.go b/tests/server/api/operator_test.go
@@ -70,7 +70,7 @@ func (suite *operatorTestSuite) TestAddRemovePeer() {
 
 func (suite *operatorTestSuite) checkAddRemovePeer(cluster *tests.TestCluster) {
 	re := suite.Require()
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	stores := []*metapb.Store{
 		{
 			Id:            1,
@@ -205,7 +205,7 @@ func (suite *operatorTestSuite) checkMergeRegionOperator(cluster *tests.TestClus
 		tests.MustPutStore(re, cluster, store)
 	}
 
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	r1 := core.NewTestRegionInfo(10, 1, []byte(""), []byte("b"), core.SetWrittenBytes(1000), core.SetReadBytes(1000), core.SetRegionConfVer(1), core.SetRegionVersion(1))
 	tests.MustPutRegionInfo(re, cluster, r1)
 	r2 := core.NewTestRegionInfo(20, 1, []byte("b"), []byte("c"), core.SetWrittenBytes(2000), core.SetReadBytes(0), core.SetRegionConfVer(2), core.SetRegionVersion(3))
@@ -241,7 +241,7 @@ func (suite *operatorTestSuite) TestTransferRegionWithPlacementRule() {
 
 func (suite *operatorTestSuite) checkTransferRegionWithPlacementRule(cluster *tests.TestCluster) {
 	re := suite.Require()
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	stores := []*metapb.Store{
 		{
 			Id:            1,
@@ -521,7 +521,7 @@ func (suite *operatorTestSuite) TestGetOperatorsAsObject() {
 
 func (suite *operatorTestSuite) checkGetOperatorsAsObject(cluster *tests.TestCluster) {
 	re := suite.Require()
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	stores := []*metapb.Store{
 		{
 			Id:            1,
@@ -642,7 +642,7 @@ func (suite *operatorTestSuite) checkRemoveOperators(cluster *tests.TestCluster)
 		tests.MustPutStore(re, cluster, store)
 	}
 
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	r1 := core.NewTestRegionInfo(10, 1, []byte(""), []byte("b"), core.SetWrittenBytes(1000), core.SetReadBytes(1000), core.SetRegionConfVer(1), core.SetRegionVersion(1))
 	tests.MustPutRegionInfo(re, cluster, r1)
 	r2 := core.NewTestRegionInfo(20, 1, []byte("b"), []byte("c"), core.SetWrittenBytes(2000), core.SetReadBytes(0), core.SetRegionConfVer(2), core.SetRegionVersion(3))

diff --git a/tests/server/api/region_test.go b/tests/server/api/region_test.go
@@ -263,7 +263,7 @@ func (suite *regionTestSuite) TestCheckRegionsReplicated() {
 
 func (suite *regionTestSuite) checkRegionsReplicated(cluster *tests.TestCluster) {
 	re := suite.Require()
-	pauseRuleChecker(re, cluster)
+	pauseAllCheckers(re, cluster)
 	leader := cluster.GetLeaderServer()
 	urlPrefix := leader.GetAddr() + "/pd/api/v1"
 
@@ -416,15 +416,16 @@ func checkRegionCount(re *require.Assertions, cluster *tests.TestCluster, count
 	}
 }
 
-// pauseRuleChecker will pause rule checker to avoid unexpected operator.
-func pauseRuleChecker(re *require.Assertions, cluster *tests.TestCluster) {
-	checkerName := "rule"
+func pauseAllCheckers(re *require.Assertions, cluster *tests.TestCluster) {
+	checkerNames := []string{"learner", "replica", "rule", "split", "merge", "joint-state"}
 	addr := cluster.GetLeaderServer().GetAddr()
-	resp := make(map[string]any)
-	url := fmt.Sprintf("%s/pd/api/v1/checker/%s", addr, checkerName)
-	err := tu.CheckPostJSON(testDialClient, url, []byte(`{"delay":1000}`), tu.StatusOK(re))
-	re.NoError(err)
-	err = tu.ReadGetJSON(re, testDialClient, url, &resp)
-	re.NoError(err)
-	re.True(resp["paused"].(bool))
+	for _, checkerName := range checkerNames {
+		resp := make(map[string]any)
+		url := fmt.Sprintf("%s/pd/api/v1/checker/%s", addr, checkerName)
+		err := tu.CheckPostJSON(testDialClient, url, []byte(`{"delay":1000}`), tu.StatusOK(re))
+		re.NoError(err)
+		err = tu.ReadGetJSON(re, testDialClient, url, &resp)
+		re.NoError(err)
+		re.True(resp["paused"].(bool))
+	}
 }