diff --git a/pkgs/strbyteconv/strbyteconv.go b/pkgs/strbyteconv/strbyteconv.go index 64a12a76d..58bd1d329 100644 --- a/pkgs/strbyteconv/strbyteconv.go +++ b/pkgs/strbyteconv/strbyteconv.go @@ -59,14 +59,20 @@ func BytesToString(bytes uint64) string { func StringToBytes(s string) (uint64, error) { var bytes uint64 + var err error - StringPattern, err := regexp.Compile(`(?i)^(-?\d+)([KMGT])B?$`) + bytes, err = strconv.ParseUint(s, 10, 64) + if err == nil { + return bytes, nil + } + + stringPattern, err := regexp.Compile(`(?i)^(-?\d+)([BKMGT])B?$`) if err != nil { return 0, err } - parts := StringPattern.FindStringSubmatch(strings.TrimSpace(s)) - if len(parts) < 3 { + parts := stringPattern.FindStringSubmatch(strings.TrimSpace(s)) + if len(parts) < 2 { return 0, errors.New("Incorrect string format must be K,KB,M,MB,G,GB") } @@ -87,6 +93,8 @@ func StringToBytes(s string) (uint64, error) { bytes = value * UNIT_KILOBYTE case "B": bytes = value * UNIT_BYTE + default: + return 0, errors.New("Incorrect string format must be K,KB,M,MB,G,GB") } return bytes, nil diff --git a/pkgs/strbyteconv/strbyteconv_test.go b/pkgs/strbyteconv/strbyteconv_test.go index 2f807ff50..6ae2e0ac1 100644 --- a/pkgs/strbyteconv/strbyteconv_test.go +++ b/pkgs/strbyteconv/strbyteconv_test.go @@ -17,8 +17,10 @@ package strbyteconv import ( - . "gopkg.in/check.v1" + "log" "testing" + + . "gopkg.in/check.v1" ) type MySuite struct{} @@ -43,7 +45,20 @@ func (s *MySuite) Test(c *C) { value = BytesToString(100 * UNIT_TERABYTE) c.Assert(value, Equals, "100TB") - bytes, err := StringToBytes("100KB") + bytes, err := StringToBytes("100B") + log.Println(err) + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100)) + + bytes, err = StringToBytes("100") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100)) + + bytes, err = StringToBytes("100KB") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100*UNIT_KILOBYTE)) + + bytes, err = StringToBytes("100K") c.Assert(err, IsNil) c.Assert(bytes, Equals, uint64(100*UNIT_KILOBYTE)) @@ -51,12 +66,49 @@ func (s *MySuite) Test(c *C) { c.Assert(err, IsNil) c.Assert(bytes, Equals, uint64(100*UNIT_MEGABYTE)) + bytes, err = StringToBytes("100M") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100*UNIT_MEGABYTE)) + bytes, err = StringToBytes("100GB") c.Assert(err, IsNil) c.Assert(bytes, Equals, uint64(100*UNIT_GIGABYTE)) + bytes, err = StringToBytes("100G") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100*UNIT_GIGABYTE)) + bytes, err = StringToBytes("100TB") c.Assert(err, IsNil) c.Assert(bytes, Equals, uint64(100*UNIT_TERABYTE)) + bytes, err = StringToBytes("100T") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(100*UNIT_TERABYTE)) + + bytes, err = StringToBytes("0") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(0)) + + bytes, err = StringToBytes("23") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(23)) + + bytes, err = StringToBytes("0TB") + c.Assert(err, IsNil) + c.Assert(bytes, Equals, uint64(0)) +} + +func (s *MySuite) TestBadInput(c *C) { + _, err := StringToBytes("") + c.Assert(err, Not(IsNil)) + + _, err = StringToBytes("HELLO") + c.Assert(err, Not(IsNil)) + + _, err = StringToBytes("-20B") + c.Assert(err, Not(IsNil)) + + _, err = StringToBytes("-20MB") + c.Assert(err, Not(IsNil)) }