From 92cc67a404cc8b1b9815602e51ec8188a1d36359 Mon Sep 17 00:00:00 2001
From: GLYASAI <chaoschunqfeijg@163.com>
Date: Sun, 11 Aug 2019 18:42:26 +0800
Subject: [PATCH] [FIX] wrong memory

---
 builder/parser/parser.go      | 45 +++++++++++++++++++++++++++++++----
 builder/parser/parser_test.go | 24 ++++++++++++++++---
 2 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/builder/parser/parser.go b/builder/parser/parser.go
index 87707008d..65cffd5cc 100644
--- a/builder/parser/parser.go
+++ b/builder/parser/parser.go
@@ -239,15 +239,50 @@ func DetermineDeployType(imageName Image) string {
 //10k 128
 //10b 128
 func readmemory(s string) int {
-	q, err := resource.ParseQuantity(strings.ToUpper(s))
+	def := 512
+	s = strings.ToLower(s)
+	// <binarySI>        ::= Ki | Mi | Gi | Ti | Pi | Ei
+	isValid := false
+	validUnits := map[string]string{
+		"gi": "Gi", "mi": "Mi", "ki": "Ki",
+	}
+	for k, v := range validUnits {
+		if strings.Contains(s, k) {
+			isValid = true
+			s = strings.Replace(s, k, v, 1)
+			break
+		}
+	}
+	if !isValid {
+		validUnits := map[string]string{
+			"g": "Gi", "m": "Mi", "k": "Ki",
+		}
+		for k, v := range validUnits {
+			if strings.Contains(s, k) {
+				isValid = true
+				s = strings.Replace(s, k, v, 1)
+				break
+			}
+		}
+	}
+	if !isValid {
+		logrus.Warningf("s: %s; invalid unit", s)
+		return def
+	}
+	q, err := resource.ParseQuantity(s)
 	if err != nil {
-		return 512
+		logrus.Warningf("s: %s; failed to parse quantity: %v", s, err)
+		return def
+	}
+	re, ok := q.AsInt64()
+	if !ok {
+		logrus.Warningf("failed to int64: %d", re)
+		return def
 	}
-	re, _ := q.AsInt64()
 	if re != 0 {
-		return int(re) / (1000 * 1000)
+		return int(re) / (1024 * 1024)
 	}
-	return 512
+	return def
 }
 
 //ParseImageName parse image name
diff --git a/builder/parser/parser_test.go b/builder/parser/parser_test.go
index c98702a3d..be917900f 100644
--- a/builder/parser/parser_test.go
+++ b/builder/parser/parser_test.go
@@ -50,7 +50,25 @@ func TestDetermineDeployType(t *testing.T) {
 }
 
 func TestReadmemory(t *testing.T) {
-	t.Log(readmemory("10G"))
-	t.Log(readmemory("300m"))
-	t.Log(readmemory("300M"))
+	testcases := []struct {
+		mem string
+		exp int
+	}{
+		{mem: "", exp: 512},
+		{mem: "2Gi", exp: 2 * 1024},
+		{mem: "2G", exp: 2 * 1024},
+		{mem: "300Mi", exp: 300},
+		{mem: "300m", exp: 300},
+		{mem: "1024Ki", exp: 1024 / 1024},
+		{mem: "1024k", exp: 1024 / 1024},
+		{mem: "1024K", exp: 1024 / 1024},
+		{mem: "1048576Bi", exp: 512},
+		{mem: "abc", exp: 512},
+	}
+	for _, tc := range testcases {
+		mem := readmemory(tc.mem)
+		if mem != tc.exp {
+			t.Errorf("mem: %s; Expected %d, but returned %d", tc.mem, tc.exp, mem)
+		}
+	}
 }