pd-ctl: fix output (#1765) (#1763) (cherry pick to release-3.1) (#1771)

Signed-off-by: lhy1024 <admin@liudos.us>

tests/pdctl/cluster/cluster_test.go

@@ -15,9 +15,6 @@ package cluster_test
 
 import (
 	"encoding/json"
-	"io/ioutil"
-	"os"
-	"path/filepath"
 	"strings"
 	"testing"
 
@@ -26,7 +23,6 @@ import (
 	"github.com/pingcap/pd/server"
 	"github.com/pingcap/pd/tests"
 	"github.com/pingcap/pd/tests/pdctl"
-	ctl "github.com/pingcap/pd/tools/pd-ctl/pdctl"
 )
 
 func Test(t *testing.T) {
@@ -50,6 +46,8 @@ func (s *clusterTestSuite) TestClusterAndPing(c *C) {
 	c.Assert(err, IsNil)
 	cluster.WaitLeader()
 	pdAddr := cluster.GetConfig().GetClientURLs()
+	i := strings.Index(pdAddr, "//")
+	pdAddr = pdAddr[i+2:]
 	cmd := pdctl.InitCommand()
 	defer cluster.Destroy()
 
@@ -61,15 +59,8 @@ func (s *clusterTestSuite) TestClusterAndPing(c *C) {
 	c.Assert(json.Unmarshal(output, ci), IsNil)
 	c.Assert(ci, DeepEquals, cluster.GetCluster())
 
-	fname := filepath.Join(os.TempDir(), "stdout")
-	old := os.Stdout
-	temp, _ := os.Create(fname)
-	os.Stdout = temp
-	ctl.Start([]string{"-u", pdAddr, "--cacert=ca.pem", "cluster"})
-	temp.Close()
-	os.Stdout = old
-	out, _ := ioutil.ReadFile(fname)
-	c.Assert(strings.Contains(string(out), "no such file or directory"), IsTrue)
+	echo := pdctl.GetEcho([]string{"-u", pdAddr, "--cacert=ca.pem", "cluster"})
+	c.Assert(strings.Contains(echo, "no such file or directory"), IsTrue)
 
 	// cluster status
 	args = []string{"-u", pdAddr, "cluster", "status"}

tests/pdctl/helper.go

@@ -17,6 +17,9 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 	"sort"
 
 	"github.com/pingcap/check"
@@ -27,6 +30,7 @@ import (
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/tests"
 	"github.com/pingcap/pd/tools/pd-ctl/pdctl"
+	ctl "github.com/pingcap/pd/tools/pd-ctl/pdctl"
 	"github.com/pingcap/pd/tools/pd-ctl/pdctl/command"
 	"github.com/spf13/cobra"
 )
@@ -132,3 +136,17 @@ func MustPutRegion(c *check.C, cluster *tests.TestCluster, regionID, storeID uin
 	c.Assert(err, check.IsNil)
 	return r
 }
+
+// GetEcho is used to get echo from stdout.
+func GetEcho(args []string) string {
+	filename := filepath.Join(os.TempDir(), "stdout")
+	old := os.Stdout
+	temp, _ := os.Create(filename)
+	os.Stdout = temp
+	ctl.Start(args)
+	temp.Close()
+	os.Stdout = old
+	out, _ := ioutil.ReadFile(filename)
+	_ = os.Remove(filename)
+	return string(out)
+}

tests/pdctl/operator/operator_test.go

@@ -194,4 +194,14 @@ func (s *operatorTestSuite) TestOperator(c *C) {
 	_, output, err = pdctl.ExecuteCommandC(cmd, args...)
 	c.Assert(err, IsNil)
 	c.Assert(strings.Contains(string(output), "scatter-region"), IsTrue)
+
+	// test echo
+	echo := pdctl.GetEcho([]string{"-u", pdAddr, "operator", "remove", "1"})
+	c.Assert(strings.Contains(echo, "Success!"), IsTrue)
+	echo = pdctl.GetEcho([]string{"-u", pdAddr, "operator", "add", "scatter-region", "1"})
+	c.Assert(strings.Contains(echo, "Success!"), IsTrue)
+	echo = pdctl.GetEcho([]string{"-u", pdAddr, "operator", "remove", "1"})
+	c.Assert(strings.Contains(echo, "Success!"), IsTrue)
+	echo = pdctl.GetEcho([]string{"-u", pdAddr, "operator", "remove", "1"})
+	c.Assert(strings.Contains(echo, "Success!"), IsFalse)
 }

tests/pdctl/region/region_test.go

@@ -15,9 +15,6 @@ package region_test
 
 import (
 	"encoding/json"
-	"io/ioutil"
-	"os"
-	"path/filepath"
 	"strings"
 	"testing"
 
@@ -29,7 +26,6 @@ import (
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/tests"
 	"github.com/pingcap/pd/tests/pdctl"
-	ctl "github.com/pingcap/pd/tools/pd-ctl/pdctl"
 )
 
 func Test(t *testing.T) {
@@ -60,15 +56,9 @@ func (s *regionTestSuite) TestRegionKeyFormat(c *C) {
 	leaderServer := cluster.GetServer(cluster.GetLeader())
 	c.Assert(leaderServer.BootstrapCluster(), IsNil)
 	pdctl.MustPutStore(c, leaderServer.GetServer(), store.Id, store.State, store.Labels)
-	fname := filepath.Join(os.TempDir(), "stdout")
-	old := os.Stdout
-	temp, _ := os.Create(fname)
-	os.Stdout = temp
-	ctl.Start([]string{"-u", url, "region", "key", "--format=raw", " "})
-	temp.Close()
-	os.Stdout = old
-	out, _ := ioutil.ReadFile(fname)
-	c.Assert(strings.Contains(string(out), "unknown flag"), IsFalse)
+
+	echo := pdctl.GetEcho([]string{"-u", url, "region", "key", "--format=raw", " "})
+	c.Assert(strings.Contains(string(echo), "unknown flag"), IsFalse)
 }
 
 func (s *regionTestSuite) TestRegion(c *C) {

tests/pdctl/scheduler/scheduler_test.go

@@ -15,6 +15,7 @@ package scheduler_test
 
 import (
 	"encoding/json"
+	"strings"
 	"testing"
 	"time"
 
@@ -131,4 +132,12 @@ func (s *schedulerTestSuite) TestScheduler(c *C) {
 	for _, scheduler := range schedulers {
 		c.Assert(expected[scheduler], Equals, true)
 	}
+
+	// test echo
+	echo := pdctl.GetEcho([]string{"-u", pdAddr, "scheduler", "add", "balance-region-scheduler"})
+	c.Assert(strings.Contains(echo, "Success!"), IsTrue)
+	echo = pdctl.GetEcho([]string{"-u", pdAddr, "scheduler", "remove", "balance-region-scheduler"})
+	c.Assert(strings.Contains(echo, "Success!"), IsTrue)
+	echo = pdctl.GetEcho([]string{"-u", pdAddr, "scheduler", "remove", "balance-region-scheduler"})
+	c.Assert(strings.Contains(echo, "Success!"), IsFalse)
 }

tools/pd-ctl/main.go

@@ -34,19 +34,21 @@ var (
 	detach   bool
 	interact bool
 	version  bool
+	help     bool
 	caPath   string
 	certPath string
 	keyPath  string
 )
 
 func init() {
-	flag.StringVarP(&url, "pd", "u", "http://127.0.0.1:2379", "The pd address")
-	flag.BoolVarP(&detach, "detach", "d", true, "Run pdctl without readline")
-	flag.BoolVarP(&interact, "interact", "i", false, "Run pdctl with readline")
-	flag.BoolVarP(&version, "version", "V", false, "print version information and exit")
-	flag.StringVar(&caPath, "cacert", "", "path of file that contains list of trusted SSL CAs.")
-	flag.StringVar(&certPath, "cert", "", "path of file that contains X509 certificate in PEM format.")
-	flag.StringVar(&keyPath, "key", "", "path of file that contains X509 key in PEM format.")
+	flag.StringVarP(&url, "pd", "u", "http://127.0.0.1:2379", "The pd address.")
+	flag.BoolVarP(&detach, "detach", "d", true, "Run pdctl without readline.")
+	flag.BoolVarP(&interact, "interact", "i", false, "Run pdctl with readline.")
+	flag.BoolVarP(&version, "version", "V", false, "Print version information and exit.")
+	flag.StringVar(&caPath, "cacert", "", "The path of file that contains list of trusted SSL CAs.")
+	flag.StringVar(&certPath, "cert", "", "The path of file that contains X509 certificate in PEM format.")
+	flag.StringVar(&keyPath, "key", "", "The path of file that contains X509 key in PEM format.")
+	flag.BoolVarP(&help, "help", "h", false, "Help message.")
 }
 
 func main() {
@@ -57,6 +59,10 @@ func main() {
 	flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true
 	flag.Parse()
 
+	if help {
+		flag.Usage()
+		os.Exit(0)
+	}
 	if version {
 		server.PrintPDInfo()
 		os.Exit(0)

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

@@ -146,7 +146,7 @@ func postJSON(cmd *cobra.Command, prefix string, input map[string]interface{}) {
 			cmd.Println(err)
 		}
 	} else {
-		cmd.Println("success!")
+		cmd.Println("Success!")
 	}
 }
 

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

@@ -408,6 +408,7 @@ func removeOperatorCommandFunc(cmd *cobra.Command, args []string) {
 		cmd.Println(err)
 		return
 	}
+	cmd.Println("Success!")
 }
 
 func parseUint64s(args []string) ([]uint64, error) {

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

@@ -315,4 +315,5 @@ func removeSchedulerCommandFunc(cmd *cobra.Command, args []string) {
 		cmd.Println(err)
 		return
 	}
+	cmd.Println("Success!")
 }