add some unit tests and some code optimizes

2 files changed, 115 insertions, 38 deletions

diff --git a/weed/shell/command_s3_circuitbreaker.go b/weed/shell/command_s3_circuitbreaker.go
index ec1cd033b..f1a8b08bb 100644
--- a/weed/shell/command_s3_circuitbreaker.go
+++ b/weed/shell/command_s3_circuitbreaker.go
@@ -5,9 +5,9 @@ import (
 	"flag"
 	"fmt"
 	"github.com/alecthomas/units"
-	"github.com/chrislusf/seaweedfs/weed/config"
 	"github.com/chrislusf/seaweedfs/weed/filer"
 	"github.com/chrislusf/seaweedfs/weed/pb/s3_pb"
+	"github.com/chrislusf/seaweedfs/weed/s3api/s3_config"
 	"io"
 	"strconv"
 	"strings"
@@ -15,6 +15,8 @@ import (
 	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
 )
 
+var LoadConfig = loadConfig
+
 func init() {
 	Commands = append(Commands, &commandS3CircuitBreaker{})
 }
@@ -23,7 +25,7 @@ type commandS3CircuitBreaker struct {
 }
 
 func (c *commandS3CircuitBreaker) Name() string {
-	return "s3.circuit.breaker"
+	return "s3.circuitBreaker"
 }
 
 func (c *commandS3CircuitBreaker) Help() string {
@@ -31,42 +33,42 @@ func (c *commandS3CircuitBreaker) Help() string {
 
 	# examples
 	# add
-	s3.circuit.breaker -actions Read,Write -values 500,200 -global -enable -apply -type count
-	s3.circuit.breaker -actions Write -values 200MiB -global -enable -apply -type bytes
-	s3.circuit.breaker -actions Write -values 200MiB -bucket x,y,z -enable -apply -type bytes
+	s3.circuitBreaker -actions Read,Write -values 500,200 -global -enable -apply -type count
+	s3.circuitBreaker -actions Write -values 200MiB -global -enable -apply -type bytes
+	s3.circuitBreaker -actions Write -values 200MiB -bucket x,y,z -enable -apply -type bytes
 
 	#delete
-	s3.circuit.breaker -actions Write -bucket x,y,z -delete -apply -type bytes
-	s3.circuit.breaker -actions Write -bucket x,y,z -delete -apply
-	s3.circuit.breaker -actions Write -delete -apply
+	s3.circuitBreaker -actions Write -bucket x,y,z -delete -apply -type bytes
+	s3.circuitBreaker -actions Write -bucket x,y,z -delete -apply
+	s3.circuitBreaker -actions Write -delete -apply
 	`
 }
 
 func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
-	dir := config.CircuitBreakerConfigDir
-	file := config.CircuitBreakerConfigFile
+	dir := s3_config.CircuitBreakerConfigDir
+	file := s3_config.CircuitBreakerConfigFile
 
 	s3CircuitBreakerCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
-	buckets := s3CircuitBreakerCommand.String("buckets", "", "comma separated buckets names")
-	global := s3CircuitBreakerCommand.Bool("global", false, "comma separated buckets names")
+	buckets := s3CircuitBreakerCommand.String("buckets", "", "the bucket name(s) to configure, eg: -buckets x,y,z")
+	global := s3CircuitBreakerCommand.Bool("global", false, "configure global circuit breaker")
 
 	actions := s3CircuitBreakerCommand.String("actions", "", "comma separated actions names: Read,Write,List,Tagging,Admin")
 	limitType := s3CircuitBreakerCommand.String("type", "", "count|bytes simultaneous requests count")
 	values := s3CircuitBreakerCommand.String("values", "", "comma separated max values,Maximum number of simultaneous requests content length, support byte unit: eg: 1k, 10m, 1g")
 
-	enabled := s3CircuitBreakerCommand.Bool("enable", true, "enable or disable circuit breaker")
-	deleted := s3CircuitBreakerCommand.Bool("delete", false, "delete users, actions or access keys")
+	disabled := s3CircuitBreakerCommand.Bool("disable", false, "disable global or buckets circuit breaker")
+	deleted := s3CircuitBreakerCommand.Bool("delete", false, "delete circuit breaker config")
 
 	apply := s3CircuitBreakerCommand.Bool("apply", false, "update and apply current configuration")
 
 	if err = s3CircuitBreakerCommand.Parse(args); err != nil {
 		return nil
+
 	}
 
 	var buf bytes.Buffer
-	if err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
-		return filer.ReadEntry(commandEnv.MasterClient, client, dir, file, &buf)
-	}); err != nil && err != filer_pb.ErrNotFound {
+	err = LoadConfig(commandEnv, dir, file, &buf)
+	if err != nil {
 		return err
 	}
 
@@ -90,7 +92,7 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 				deleteGlobalActions(cbCfg, cmdActions, limitType)
 				if cbCfg.Buckets != nil {
 					var allBuckets []string
-					for bucket, _ := range cbCfg.Buckets {
+					for bucket := range cbCfg.Buckets {
 						allBuckets = append(allBuckets, bucket)
 					}
 					deleteBucketsActions(allBuckets, cbCfg, cmdActions, limitType)
@@ -108,7 +110,7 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 			}
 		}
 	} else {
-		cmdBuckets, cmdActions, cmdValues, err := c.initActionsAndValues(buckets, actions, limitType, values, false)
+		cmdBuckets, cmdActions, cmdValues, err := c.initActionsAndValues(buckets, actions, limitType, values, *disabled)
 		if err != nil {
 			return err
 		}
@@ -125,7 +127,7 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 					cbOptions = &s3_pb.CbOptions{}
 					cbCfg.Buckets[bucket] = cbOptions
 				}
-				cbOptions.Enabled = *enabled
+				cbOptions.Enabled = !*disabled
 
 				if len(cmdActions) > 0 {
 					err = insertOrUpdateValues(cbOptions, cmdActions, cmdValues, limitType)
@@ -146,7 +148,7 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 				globalOptions = &s3_pb.CbOptions{Actions: make(map[string]int64, len(cmdActions))}
 				cbCfg.Global = globalOptions
 			}
-			globalOptions.Enabled = *enabled
+			globalOptions.Enabled = !*disabled
 
 			if len(cmdActions) > 0 {
 				err = insertOrUpdateValues(globalOptions, cmdActions, cmdValues, limitType)
@@ -167,8 +169,8 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 		return err
 	}
 
-	fmt.Fprintf(writer, string(buf.Bytes()))
-	fmt.Fprintln(writer)
+	_, _ = fmt.Fprintf(writer, string(buf.Bytes()))
+	_, _ = fmt.Fprintln(writer)
 
 	if *apply {
 		if err := commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
@@ -176,12 +178,20 @@ func (c *commandS3CircuitBreaker) Do(args []string, commandEnv *CommandEnv, writ
 		}); err != nil {
 			return err
 		}
-
 	}
 
 	return nil
 }
 
+func loadConfig(commandEnv *CommandEnv, dir string, file string, buf *bytes.Buffer) error {
+	if err := commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
+		return filer.ReadEntry(commandEnv.MasterClient, client, dir, file, buf)
+	}); err != nil && err != filer_pb.ErrNotFound {
+		return err
+	}
+	return nil
+}
+
 func insertOrUpdateValues(cbOptions *s3_pb.CbOptions, cmdActions []string, cmdValues []int64, limitType *string) error {
 	if len(*limitType) == 0 {
 		return fmt.Errorf("type not valid, only 'count' and 'bytes' are allowed")
@@ -193,7 +203,7 @@ func insertOrUpdateValues(cbOptions *s3_pb.CbOptions, cmdActions []string, cmdVa
 
 	if len(cmdValues) > 0 {
 		for i, action := range cmdActions {
-			cbOptions.Actions[config.Concat(action, *limitType)] = cmdValues[i]
+			cbOptions.Actions[s3_config.Concat(action, *limitType)] = cmdValues[i]
 		}
 	}
 	return nil
@@ -213,7 +223,7 @@ func deleteBucketsActions(cmdBuckets []string, cbCfg *s3_pb.S3CircuitBreakerConf
 			if cbOption, ok := cbCfg.Buckets[bucket]; ok {
 				if len(cmdActions) > 0 && cbOption.Actions != nil {
 					for _, action := range cmdActions {
-						delete(cbOption.Actions, config.Concat(action, *limitType))
+						delete(cbOption.Actions, s3_config.Concat(action, *limitType))
 					}
 				}
 
@@ -240,7 +250,7 @@ func deleteGlobalActions(cbCfg *s3_pb.S3CircuitBreakerConfig, cmdActions []strin
 		return
 	} else {
 		for _, action := range cmdActions {
-			delete(globalOptions.Actions, config.Concat(action, *limitType))
+			delete(globalOptions.Actions, s3_config.Concat(action, *limitType))
 		}
 	}
 
@@ -249,7 +259,7 @@ func deleteGlobalActions(cbCfg *s3_pb.S3CircuitBreakerConfig, cmdActions []strin
 	}
 }
 
-func (c *commandS3CircuitBreaker) initActionsAndValues(buckets, actions, limitType, values *string, deleteOp bool) (cmdBuckets, cmdActions []string, cmdValues []int64, err error) {
+func (c *commandS3CircuitBreaker) initActionsAndValues(buckets, actions, limitType, values *string, parseValues bool) (cmdBuckets, cmdActions []string, cmdValues []int64, err error) {
 	if len(*buckets) > 0 {
 		cmdBuckets = strings.Split(*buckets, ",")
 	}
@@ -260,27 +270,27 @@ func (c *commandS3CircuitBreaker) initActionsAndValues(buckets, actions, limitTy
 		//check action valid
 		for _, action := range cmdActions {
 			var found bool
-			for _, allowedAction := range config.AllowedActions {
+			for _, allowedAction := range s3_config.AllowedActions {
 				if allowedAction == action {
 					found = true
 				}
 			}
 			if !found {
-				return nil, nil, nil, fmt.Errorf("value(%s) of flag[-action] not valid, allowed actions: %v", *actions, config.AllowedActions)
+				return nil, nil, nil, fmt.Errorf("value(%s) of flag[-action] not valid, allowed actions: %v", *actions, s3_config.AllowedActions)
 			}
 		}
 	}
 
-	if !deleteOp {
+	if !parseValues {
 		if len(cmdActions) < 0 {
-			for _, action := range config.AllowedActions {
+			for _, action := range s3_config.AllowedActions {
 				cmdActions = append(cmdActions, action)
 			}
 		}
 
 		if len(*limitType) > 0 {
 			switch *limitType {
-			case config.LimitTypeCount:
+			case s3_config.LimitTypeCount:
 				elements := strings.Split(*values, ",")
 				if len(cmdActions) != len(elements) {
 					if len(elements) != 1 || len(elements) == 0 {
@@ -288,7 +298,7 @@ func (c *commandS3CircuitBreaker) initActionsAndValues(buckets, actions, limitTy
 					}
 					v, err := strconv.Atoi(elements[0])
 					if err != nil {
-						return nil, nil, nil, fmt.Errorf("value of -counts must be a legal number(s)")
+						return nil, nil, nil, fmt.Errorf("value of -values must be a legal number(s)")
 					}
 					for range cmdActions {
 						cmdValues = append(cmdValues, int64(v))
@@ -297,16 +307,16 @@ func (c *commandS3CircuitBreaker) initActionsAndValues(buckets, actions, limitTy
 					for _, value := range elements {
 						v, err := strconv.Atoi(value)
 						if err != nil {
-							return nil, nil, nil, fmt.Errorf("value of -counts must be a legal number(s)")
+							return nil, nil, nil, fmt.Errorf("value of -values must be a legal number(s)")
 						}
 						cmdValues = append(cmdValues, int64(v))
 					}
 				}
-			case config.LimitTypeBytes:
+			case s3_config.LimitTypeBytes:
 				elements := strings.Split(*values, ",")
 				if len(cmdActions) != len(elements) {
 					if len(elements) != 1 || len(elements) == 0 {
-						return nil, nil, nil, fmt.Errorf("count of flag[-actions] and flag[-counts] not equal")
+						return nil, nil, nil, fmt.Errorf("values count of -actions and -values not equal")
 					}
 					v, err := units.ParseStrictBytes(elements[0])
 					if err != nil {
diff --git a/weed/shell/command_s3_circuitbreaker_test.go b/weed/shell/command_s3_circuitbreaker_test.go
index 93e9fa414..214256b5c 100644
--- a/weed/shell/command_s3_circuitbreaker_test.go
+++ b/weed/shell/command_s3_circuitbreaker_test.go
@@ -1,7 +1,74 @@
 package shell
 
-import "testing"
+import (
+	"bytes"
+	"strings"
+	"testing"
+)
+
+type Case struct {
+	args   []string
+	result string
+}
+
+var (
+	TestCases = []*Case{
+		//add circuit breaker config for global
+		{
+			args:   strings.Split("-global -type count -actions Read,Write -values 500,200", " "),
+			result: "{\n  \"global\": {\n    \"enabled\": true,\n    \"actions\": {\n      \"Read:count\": \"500\",\n      \"Write:count\": \"200\"\n    }\n  }\n}\n",
+		},
+		//disable global config
+		{
+			args:   strings.Split("-global -disable", " "),
+			result: "{\n  \"global\": {\n    \"actions\": {\n      \"Read:count\": \"500\",\n      \"Write:count\": \"200\"\n    }\n  }\n}\n",
+		},
+		//add circuit breaker config for buckets x,y,z
+		{
+			args:   strings.Split("-buckets x,y,z -type count -actions Read,Write -values 200,100", " "),
+			result: "{\n  \"global\": {\n    \"actions\": {\n      \"Read:count\": \"500\",\n      \"Write:count\": \"200\"\n    }\n  },\n  \"buckets\": {\n    \"x\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    },\n    \"y\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    },\n    \"z\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    }\n  }\n}\n",
+		},
+		//disable circuit breaker config of x
+		{
+			args:   strings.Split("-buckets x -disable", " "),
+			result: "{\n  \"global\": {\n    \"actions\": {\n      \"Read:count\": \"500\",\n      \"Write:count\": \"200\"\n    }\n  },\n  \"buckets\": {\n    \"x\": {\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    },\n    \"y\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    },\n    \"z\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    }\n  }\n}\n",
+		},
+		//delete circuit breaker config of x
+		{
+			args:   strings.Split("-buckets x -delete", " "),
+			result: "{\n  \"global\": {\n    \"actions\": {\n      \"Read:count\": \"500\",\n      \"Write:count\": \"200\"\n    }\n  },\n  \"buckets\": {\n    \"y\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    },\n    \"z\": {\n      \"enabled\": true,\n      \"actions\": {\n        \"Read:count\": \"200\",\n        \"Write:count\": \"100\"\n      }\n    }\n  }\n}\n",
+		},
+		//clear all circuit breaker config
+		{
+			args:   strings.Split("-delete", " "),
+			result: "{\n\n}\n",
+		},
+	}
+)
 
 func TestCircuitBreakerShell(t *testing.T) {
+	var writeBuf bytes.Buffer
+	cmd := &commandS3CircuitBreaker{}
+	LoadConfig = func(commandEnv *CommandEnv, dir string, file string, buf *bytes.Buffer) error {
+		_, err := buf.Write(writeBuf.Bytes())
+		if err != nil {
+			return err
+		}
+		writeBuf.Reset()
+		return nil
+	}
+
+	for i, tc := range TestCases {
+		err := cmd.Do(tc.args, nil, &writeBuf)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if i != 0 {
+			result := writeBuf.String()
+			if result != tc.result {
+				t.Fatal("result of s3 circuit breaker shell command is unexpect!")
+			}
 
+		}
+	}
 }