close tikv#4490: add teset

Signed-off-by: Cabinfever_B <[email protected]>

pkg/apiutil/apiutil.go

@@ -27,6 +27,14 @@ import (
 	"github.com/unrolled/render"
 )
 
+var (
+	// ComponentSignatureKey is used for http request header key
+	// to identify component signature
+	ComponentSignatureKey = "ti-component"
+	// ComponentAnonymousValue identify anonymous request source
+	ComponentAnonymousValue = "anonymous"
+)
+
 // DeferClose captures the error returned from closing (if an error occurs).
 // This is designed to be used in a defer statement.
 func DeferClose(c io.Closer, err *error) {
@@ -127,3 +135,25 @@ func ErrorResp(rd *render.Render, w http.ResponseWriter, err error) {
 		rd.JSON(w, http.StatusInternalServerError, err.Error())
 	}
 }
+
+// GetComponentNameOnHTTP return component name from Request Header
+func GetComponentNameOnHTTP(r *http.Request) string {
+	componentName := r.Header.Get(ComponentSignatureKey)
+	if componentName == "" {
+		componentName = ComponentAnonymousValue
+	}
+	return componentName
+}
+
+type UserSignatureRoundTripper struct {
+	Proxied   http.RoundTripper
+	Component string
+}
+
+// RoundTrip is used to implement RoundTripper
+func (rt *UserSignatureRoundTripper) RoundTrip(req *http.Request) (resp *http.Response, err error) {
+	req.Header.Add(ComponentSignatureKey, rt.Component)
+	// Send the request, get the response and the error
+	resp, err = rt.Proxied.RoundTrip(req)
+	return
+}

tests/pdctl/global.go → tests/pdctl/global_test.go

@@ -16,15 +16,17 @@ package pdctl
 
 import (
 	"context"
-	"encoding/json"
+	"fmt"
+	"net/http"
 	"testing"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/log"
+	"github.com/tikv/pd/pkg/apiutil"
+	"github.com/tikv/pd/pkg/testutil"
 	"github.com/tikv/pd/server"
-	"github.com/tikv/pd/server/api"
-	"github.com/tikv/pd/tests"
-	"github.com/tikv/pd/tools/pd-ctl/pdctl"
 	cmd "github.com/tikv/pd/tools/pd-ctl/pdctl"
+	"go.uber.org/zap"
 )
 
 func Test(t *testing.T) {
@@ -39,21 +41,39 @@ func (s *globalTestSuite) SetUpSuite(c *C) {
 	server.EnableZap = true
 }
 
-func (s *globalTestSuite) TestHealth(c *C) {
+func (s *globalTestSuite) TestSendAndGetComponent(c *C) {
+	handler := func(ctx context.Context, s *server.Server) (http.Handler, server.ServiceGroup, error) {
+		mux := http.NewServeMux()
+		mux.HandleFunc("/pd/api/v1/health", func(w http.ResponseWriter, r *http.Request) {
+			component := apiutil.GetComponentNameOnHTTP(r)
+			for k := range r.Header {
+				log.Info("header", zap.String("key", k))
+			}
+			log.Info("component", zap.String("component", component))
+			c.Assert(component, Equals, "pdctl")
+			fmt.Fprint(w, component)
+		})
+		info := server.ServiceGroup{
+			IsCore: true,
+		}
+		return mux, info, nil
+	}
+	cfg := server.NewTestSingleConfig(checkerWithNilAssert(c))
 	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	cluster, err := tests.NewTestCluster(ctx, 1)
+	svr, err := server.CreateServer(ctx, cfg, handler)
 	c.Assert(err, IsNil)
-	err = cluster.RunInitialServers()
+	err = svr.Run()
 	c.Assert(err, IsNil)
-	pdAddr := cluster.GetConfig().GetClientURL()
-	cmd := cmd.GetRootCmd()
+	pdAddr := svr.GetAddr()
+	defer func() {
+		cancel()
+		svr.Close()
+		testutil.CleanServer(svr.GetConfig().DataDir)
+	}()
 
+	cmd := cmd.GetRootCmd()
 	args := []string{"-u", pdAddr, "health"}
-	output, err := pdctl.ExecuteCommand(cmd, args...)
-	c.Assert(err, IsNil)
-	h := make([]api.Health, len(healths))
-	c.Assert(json.Unmarshal(output, &h), IsNil)
+	output, err := ExecuteCommand(cmd, args...)
 	c.Assert(err, IsNil)
-	c.Assert(h, DeepEquals, healths)
+	c.Assert(string(output), Equals, "pdctl\n")
 }

tests/pdctl/helper.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/spf13/cobra"
+	"github.com/tikv/pd/pkg/assertutil"
 	"github.com/tikv/pd/server"
 	"github.com/tikv/pd/server/api"
 	"github.com/tikv/pd/server/core"
@@ -113,3 +114,11 @@ func MustPutRegion(c *check.C, cluster *tests.TestCluster, regionID, storeID uin
 	c.Assert(err, check.IsNil)
 	return r
 }
+
+func checkerWithNilAssert(c *check.C) *assertutil.Checker {
+	checker := assertutil.NewChecker(c.FailNow)
+	checker.IsNil = func(obtained interface{}) {
+		c.Assert(obtained, check.IsNil)
+	}
+	return checker
+}

tools/pd-ctl/pdctl/command/global.go

@@ -25,7 +25,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/spf13/cobra"
-	"go.etcd.io/etcd/pkg/transport"
+	"github.com/tikv/pd/pkg/apiutil"
 )
 
 var (
@@ -35,19 +35,11 @@ var (
 
 // InitHTTPSClient creates https client with ca file
 func InitHTTPSClient(caPath, certPath, keyPath string) error {
-	tlsInfo := transport.TLSInfo{
-		CertFile:      certPath,
-		KeyFile:       keyPath,
-		TrustedCAFile: caPath,
-	}
-	tlsConfig, err := tlsInfo.ClientConfig()
-	if err != nil {
-		return errors.WithStack(err)
-	}
 
 	dialClient = &http.Client{
-		Transport: &http.Transport{
-			TLSClientConfig: tlsConfig,
+		Transport: &apiutil.UserSignatureRoundTripper{
+			Component: "pdctl",
+			Proxied:   http.DefaultTransport,
 		},
 	}