From 5abcaf4d588861927e76ec861aeefe7c08c43f09 Mon Sep 17 00:00:00 2001
From: "Frederick F. Kautz IV" <fkautz@alumni.cmu.edu>
Date: Thu, 4 Dec 2014 23:03:05 -0800
Subject: [PATCH] Adding more unit tests to strbyteconv. Fixes #69.

---
 pkgs/strbyteconv/strbyteconv.go      | 14 +++++--
 pkgs/strbyteconv/strbyteconv_test.go | 56 +++++++++++++++++++++++++++-
 2 files changed, 65 insertions(+), 5 deletions(-)

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))
 }