Browse Source

Merge pull request #5902 from mitake/bench-auth

tools: add --user for auth in benchmarks
Xiang Li 9 years ago
parent
commit
4bc29e2b9c
2 changed files with 16 additions and 0 deletions
  1. 4 0
      tools/benchmark/cmd/root.go
  2. 12 0
      tools/benchmark/cmd/util.go

+ 4 - 0
tools/benchmark/cmd/root.go

@@ -46,6 +46,8 @@ var (
 
 	cpuProfPath string
 	memProfPath string
+
+	user string
 )
 
 func init() {
@@ -57,4 +59,6 @@ func init() {
 	RootCmd.PersistentFlags().StringVar(&tls.CertFile, "cert", "", "identify HTTPS client using this SSL certificate file")
 	RootCmd.PersistentFlags().StringVar(&tls.KeyFile, "key", "", "identify HTTPS client using this SSL key file")
 	RootCmd.PersistentFlags().StringVar(&tls.CAFile, "cacert", "", "verify certificates of HTTPS-enabled servers using this CA bundle")
+
+	RootCmd.PersistentFlags().StringVar(&user, "user", "", "specify username and password in username:password format")
 }

+ 12 - 0
tools/benchmark/cmd/util.go

@@ -18,6 +18,7 @@ import (
 	"crypto/rand"
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/coreos/etcd/clientv3"
 )
@@ -41,6 +42,17 @@ func mustCreateConn() *clientv3.Client {
 		cfg.TLS = cfgtls
 	}
 
+	if len(user) != 0 {
+		splitted := strings.SplitN(user, ":", 2)
+		if len(splitted) != 2 {
+			fmt.Fprintf(os.Stderr, "bad user information: %s\n", user)
+			os.Exit(1)
+		}
+
+		cfg.Username = splitted[0]
+		cfg.Password = splitted[1]
+	}
+
 	client, err := clientv3.New(cfg)
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "dial error: %v\n", err)