Merge pull request #1494 from alexbozhenko/b_is_for_bytes

Allow B and b when specifying number of bytes

Author: ripienaar

cli/bench_command.go

@@ -92,7 +92,7 @@ func configureBenchCommand(app commandHost) {
 		f.Flag("msgs", "Number of messages to publish or subscribe to").Default("100000").IntVar(&c.numMsg)
 		f.Flag("progress", "Enable or disable the progress bar").Default("true").BoolVar(&c.progressBar)
 		f.Flag("csv", "Save benchmark data to CSV file").StringVar(&c.csvFile)
-		f.Flag("size", "Size of the test messages").Default("128").StringVar(&c.msgSizeString)
+		f.Flag("size", "Size of the test messages").Default("128B").StringVar(&c.msgSizeString)
 		// TODO: support randomized payload data
 	}
 

internal/util/util.go

@@ -595,7 +595,7 @@ func Base64IfNotPrintable(val []byte) string {
 }
 
 var bytesUnitSplitter = regexp.MustCompile(`^(\d+)(\w+)`)
-var errInvalidByteString = errors.New("bytes must end in K, KB, M, MB, G, GB, T or TB")
+var errInvalidByteString = errors.New("bytes must end in B, K, KB, M, MB, G, GB, T or TB")
 var errBytesOutOfRangeString = errors.New("bytes value out of allowed range")
 var errInvalidBitSize = errors.New("invalid bit size")
 var validBitSizes = []int{8, 16, 32, 64}
@@ -645,7 +645,7 @@ func ParseStringAsBytes(s string, bitSize int) (int64, error) {
 	}
 
 	suffix := matches[2]
-	suffixMap := map[string]int64{"K": 10, "KB": 10, "KIB": 10, "M": 20, "MB": 20, "MIB": 20, "G": 30, "GB": 30, "GIB": 30, "T": 40, "TB": 40, "TIB": 40}
+	suffixMap := map[string]int64{"B": 0, "K": 10, "KB": 10, "KIB": 10, "M": 20, "MB": 20, "MIB": 20, "G": 30, "GB": 30, "GIB": 30, "T": 40, "TB": 40, "TIB": 40}
 
 	mult, ok := suffixMap[strings.ToUpper(suffix)]
 	if !ok {

internal/util/util_test.go

@@ -210,13 +210,17 @@ func TestParseStringAsBytes(t *testing.T) {
 		error     bool
 		errorKind error
 	}{
+		{input: "1B", expect: 1},
+		{input: "1b", expect: 1},
 		{input: "1", expect: 1},
 		{input: "1000", expect: 1000},
 		{input: "1K", expect: 1024},
 		{input: "1k", expect: 1024},
 		{input: "1KB", expect: 1024},
 		{input: "1KiB", expect: 1024},
 		{input: "1kb", expect: 1024},
+		{input: "1024B", expect: 1024},
+		{input: "1024b", expect: 1024},
 		{input: "1M", expect: 1024 * 1024},
 		{input: "1MB", expect: 1024 * 1024},
 		{input: "1MiB", expect: 1024 * 1024},
@@ -232,10 +236,10 @@ func TestParseStringAsBytes(t *testing.T) {
 		{input: "9223372036854775807", expect: math.MaxInt64},
 		{input: "9223372036854775807", bits: 32, error: true, errorKind: errBytesOutOfRangeString},
 		{input: "9223372036854775807", bits: 12, error: true, errorKind: errInvalidBitSize},
+		{input: "", expect: -1},
 		{input: "-1", expect: -1},
 		{input: "-10", expect: -1},
 		{input: "-10GB", expect: -1},
-		{input: "1B", error: true},
 		{input: "1FOO", error: true},
 		{input: "FOO", error: true},
 	}