From a4b897482b319970fe1ea5c6133be66454881ff9 Mon Sep 17 00:00:00 2001
From: Kaviraj <kavirajkanagaraj@gmail.com>
Date: Wed, 8 Sep 2021 11:04:15 +0200
Subject: [PATCH] Add custom UnmarshalJSON for bytesize type (#4289)

* Add custom UnmarshalJSON for bytesize type

* Add more test case

* Fix limits test
---
 pkg/util/flagext/bytesize.go      | 12 ++++++++++
 pkg/util/flagext/bytesize_test.go | 37 +++++++++++++++++++++++++++++++
 pkg/validation/limits_test.go     |  2 +-
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/pkg/util/flagext/bytesize.go b/pkg/util/flagext/bytesize.go
index 76c61d69d413..0d59a6aba271 100644
--- a/pkg/util/flagext/bytesize.go
+++ b/pkg/util/flagext/bytesize.go
@@ -1,6 +1,7 @@
 package flagext
 
 import (
+	"encoding/json"
 	"strings"
 
 	"github.com/c2h5oh/datasize"
@@ -44,3 +45,14 @@ func (bs *ByteSize) UnmarshalYAML(unmarshal func(interface{}) error) error {
 
 	return bs.Set(str)
 }
+
+// UnmarshalJSON implements json.Unmarsal interface to work with JSON.
+func (bs *ByteSize) UnmarshalJSON(val []byte) error {
+	var str string
+
+	if err := json.Unmarshal(val, &str); err != nil {
+		return err
+	}
+
+	return bs.Set(str)
+}
diff --git a/pkg/util/flagext/bytesize_test.go b/pkg/util/flagext/bytesize_test.go
index 95331af97797..363645f9c1bb 100644
--- a/pkg/util/flagext/bytesize_test.go
+++ b/pkg/util/flagext/bytesize_test.go
@@ -3,6 +3,8 @@ package flagext
 import (
 	"testing"
 
+	"encoding/json"
+
 	"github.com/stretchr/testify/require"
 	"gopkg.in/yaml.v2"
 )
@@ -102,3 +104,38 @@ func Test_ByteSizeYAML(t *testing.T) {
 		})
 	}
 }
+
+func Test_ByteSizeJSON(t *testing.T) {
+	for _, tc := range []struct {
+		in  string
+		err bool
+		out ByteSize
+	}{
+		{
+			in:  `{ "bytes": "256GB" }`,
+			out: ByteSize(256 << 30),
+		},
+		{
+			// JSON shouldn't allow to set integer as value for ByteSize field.
+			in:  `{ "bytes": 2.62144e+07 }`,
+			err: true,
+		},
+		{
+			in:  `{ "bytes": "abc" }`,
+			err: true,
+		},
+	} {
+		t.Run(tc.in, func(t *testing.T) {
+			var out struct {
+				Bytes ByteSize `json:"bytes"`
+			}
+			err := json.Unmarshal([]byte(tc.in), &out)
+			if tc.err {
+				require.NotNil(t, err)
+			} else {
+				require.Nil(t, err)
+				require.Equal(t, tc.out, out.Bytes)
+			}
+		})
+	}
+}
diff --git a/pkg/validation/limits_test.go b/pkg/validation/limits_test.go
index 9605bbd44f5e..727f86b92413 100644
--- a/pkg/validation/limits_test.go
+++ b/pkg/validation/limits_test.go
@@ -76,7 +76,7 @@ per_tenant_override_period: 230s
   "reject_old_samples_max_age": "40s",
   "creation_grace_period": "50s",
   "enforce_metric_name": true,
-  "max_line_size": 60,
+  "max_line_size": "60",
   "max_line_size_truncate": true,
   "max_streams_per_user": 70,
   "max_global_streams_per_user": 80,