Add memory size parser

Author: tommysitu

core/util/memory_size.go

@@ -0,0 +1,48 @@
+package util
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+// MemorySize is a custom type for parsing memory sizes (e.g., "128KB", "2MB")
+type MemorySize int64
+
+// String returns the string representation of the memory size
+func (m *MemorySize) String() string {
+	return fmt.Sprintf("%d", *m)
+}
+
+// ToBytes returns the memory size as an int64 in bytes.
+func (m *MemorySize) ToBytes() int64 {
+	return int64(*m)
+}
+
+// Set parses a string like "128KB" or "2MB" and converts it to bytes
+func (m *MemorySize) Set(value string) error {
+	multiplier := int64(1)
+
+	value = strings.ToUpper(strings.TrimSpace(value))
+
+	switch {
+	case strings.HasSuffix(value, "KB"):
+		multiplier = 1024
+		value = strings.TrimSuffix(value, "KB")
+	case strings.HasSuffix(value, "MB"):
+		multiplier = 1024 * 1024
+		value = strings.TrimSuffix(value, "MB")
+	case strings.HasSuffix(value, "GB"):
+		multiplier = 1024 * 1024 * 1024
+		value = strings.TrimSuffix(value, "GB")
+	}
+
+	size, err := strconv.ParseInt(value, 10, 64)
+	if err != nil || size < 0 {
+		return fmt.Errorf("invalid memory size: %s", value)
+	}
+
+	*m = MemorySize(size * multiplier)
+	return nil
+}
+

core/util/memory_size_test.go

@@ -0,0 +1,68 @@
+package util
+
+import "testing"
+import . "github.com/onsi/gomega"
+
+func TestMemorySize_Set(t *testing.T) {
+	RegisterTestingT(t) // Register Gomega for the test
+
+	t.Run("valid inputs", func(t *testing.T) {
+		var ms MemorySize
+
+		// Test inputs with valid values
+		err := ms.Set("128KB")
+		Expect(err).To(BeNil())
+		Expect(ms).To(Equal(MemorySize(128 * 1024)))
+
+		err = ms.Set("2MB")
+		Expect(err).To(BeNil())
+		Expect(ms).To(Equal(MemorySize(2 * 1024 * 1024)))
+
+		err = ms.Set("1GB")
+		Expect(err).To(BeNil())
+		Expect(ms).To(Equal(MemorySize(1 * 1024 * 1024 * 1024)))
+
+		err = ms.Set("1024") // No suffix, treat as bytes
+		Expect(err).To(BeNil())
+		Expect(ms).To(Equal(MemorySize(1024)))
+
+		err = ms.Set(" 64MB ") // Test with leading/trailing spaces
+		Expect(err).To(BeNil())
+		Expect(ms).To(Equal(MemorySize(64 * 1024 * 1024)))
+	})
+
+	t.Run("invalid inputs", func(t *testing.T) {
+		var ms MemorySize
+
+		// Test inputs with invalid values
+		Expect(ms.Set("10XYZ")).To(Not(BeNil())) // Unknown unit
+		Expect(ms.Set("ABC")).To(Not(BeNil()))   // Non-numeric input
+		Expect(ms.Set("")).To(Not(BeNil()))      // Empty input
+		Expect(ms.Set("-5MB")).To(Not(BeNil()))  // Negative value
+	})
+
+	t.Run("boundary cases", func(t *testing.T) {
+		var ms MemorySize
+
+		// Test extremely large numbers
+		err := ms.Set("1099511627776GB") // 1 PB (petabyte, very large)
+		Expect(err).To(BeNil())
+
+		// Overflow handling (in practice, you'd want to handle overflow explicitly)
+		err = ms.Set("9223372036854775808") // Larger than int64 max
+		Expect(err).To(Not(BeNil()))
+	})
+}
+
+func TestMemorySize_AsBytes(t *testing.T) {
+	RegisterTestingT(t) // Register Gomega for the test
+
+	var ms MemorySize
+
+	// Set a value and check its string representation
+	ms = 128 * 1024
+	Expect(ms.ToBytes()).To(Equal(int64(131072)))
+
+	ms = 2 * 1024 * 1024
+	Expect(ms.ToBytes()).To(Equal(int64(2097152)))
+}