Fix pd-ctl echo output (tikv#1765)

Signed-off-by: lhy1024 <admin@liudos.us>
lhy1024 · Sep 24, 2019 · c973515 · c973515
1 parent 5c648dc
commit c973515
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 28 deletions.
diff --git a/tests/pdctl/cluster/cluster_test.go b/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) {
@@ -63,16 +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)
-	os.Remove(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"}

diff --git a/tests/pdctl/helper.go b/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)
+}
diff --git a/tests/pdctl/operator/operator_test.go b/tests/pdctl/operator/operator_test.go
@@ -195,4 +195,12 @@ 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", "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)
 }
diff --git a/tests/pdctl/region/region_test.go b/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,16 +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)
-	os.Remove(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) {

diff --git a/tests/pdctl/scheduler/scheduler_test.go b/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)
 }
diff --git a/tools/pd-ctl/pdctl/command/operator.go b/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) {

diff --git a/tools/pd-ctl/pdctl/command/scheduler.go b/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!")
 }