From af61029ac1afabd4e35189a826fd013821cb9e8d Mon Sep 17 00:00:00 2001
From: David Zhu <dyzz@google.com>
Date: Thu, 19 Apr 2018 16:37:21 -0700
Subject: [PATCH] Made storage shown in PV more readable

---
 pkg/controller/controller.go      | 24 +++++++++++++++--
 pkg/controller/controller_test.go | 43 +++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
index 8903ae527e..e860880f05 100644
--- a/pkg/controller/controller.go
+++ b/pkg/controller/controller.go
@@ -19,6 +19,7 @@ package controller
 import (
 	"context"
 	"fmt"
+	"math"
 	"net"
 	"strings"
 	"time"
@@ -359,7 +360,6 @@ func (p *csiProvisioner) Provision(options controller.VolumeOptions) (*v1.Persis
 		}
 		return nil, capErr
 	}
-	repBytesString := fmt.Sprintf("%v", respCap)
 	pv := &v1.PersistentVolume{
 		ObjectMeta: metav1.ObjectMeta{
 			Name: share,
@@ -368,7 +368,7 @@ func (p *csiProvisioner) Provision(options controller.VolumeOptions) (*v1.Persis
 			PersistentVolumeReclaimPolicy: options.PersistentVolumeReclaimPolicy,
 			AccessModes:                   options.PVC.Spec.AccessModes,
 			Capacity: v1.ResourceList{
-				v1.ResourceName(v1.ResourceStorage): resource.MustParse(repBytesString),
+				v1.ResourceName(v1.ResourceStorage): bytesToGiQuantity(respCap),
 			},
 			// TODO wait for CSI VolumeSource API
 			PersistentVolumeSource: v1.PersistentVolumeSource{
@@ -464,3 +464,23 @@ func getCredentialsFromSecret(k8s kubernetes.Interface, secretName, nameSpace st
 
 	return credentials, nil
 }
+
+func bytesToGiQuantity(bytes int64) resource.Quantity {
+	var num int64
+	var floatBytes, MiB, GiB float64
+	var suffix string
+	floatBytes = float64(bytes)
+	MiB = 1024 * 1024
+	GiB = MiB * 1024
+	// Need to give Quantity nice whole numbers or else it
+	// sometimes spits out the value in milibytes. We round up.
+	if floatBytes < GiB {
+		num = int64(math.Ceil(floatBytes / MiB))
+		suffix = "Mi"
+	} else {
+		num = int64(math.Ceil(floatBytes / GiB))
+		suffix = "Gi"
+	}
+	stringQuantity := fmt.Sprintf("%v%s", num, suffix)
+	return resource.MustParse(stringQuantity)
+}
diff --git a/pkg/controller/controller_test.go b/pkg/controller/controller_test.go
index ceebf23382..defbc92e53 100644
--- a/pkg/controller/controller_test.go
+++ b/pkg/controller/controller_test.go
@@ -390,6 +390,49 @@ func TestGetDriverName(t *testing.T) {
 	}
 }
 
+func TestBytesToQuantity(t *testing.T) {
+	tests := []struct {
+		testName    string
+		bytes       float64
+		quantString string
+	}{
+		{
+			"Gibibyte rounding up from above .5",
+			5.56 * 1024 * 1024 * 1024,
+			"6Gi",
+		},
+		{
+			"Gibibyte rounding up from below .5",
+			5.23 * 1024 * 1024 * 1024,
+			"6Gi",
+		},
+		{
+			"Gibibyte exact",
+			5 * 1024 * 1024 * 1024,
+			"5Gi",
+		},
+		{
+			"Mebibyte rounding up from below .5",
+			5.23 * 1024 * 1024,
+			"6Mi",
+		},
+		{
+			"Mebibyte/Gibibyte barrier (Quantity type rounds this)",
+			// (1024 * 1024 * 1024) - 1
+			1073741823,
+			"1Gi",
+		},
+	}
+
+	for _, test := range tests {
+		q := bytesToGiQuantity(int64(test.bytes))
+		if q.String() != test.quantString {
+			t.Errorf("test: %s, expected: %v, got: %v", test.testName, test.quantString, q.String())
+		}
+	}
+
+}
+
 func TestCreateDriverReturnsInvalidCapacityDuringProvision(t *testing.T) {
 	// Set up mocks
 	var requestedBytes int64 = 100