github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/flagext/bytesize_test.go (about)

     1  package flagext
     2  
     3  import (
     4  	"testing"
     5  
     6  	"encoding/json"
     7  
     8  	"github.com/stretchr/testify/require"
     9  	"gopkg.in/yaml.v2"
    10  )
    11  
    12  func Test_ByteSize(t *testing.T) {
    13  	for _, tc := range []struct {
    14  		in  string
    15  		err bool
    16  		out int
    17  	}{
    18  		{
    19  			in:  "abc",
    20  			err: true,
    21  		},
    22  		{
    23  			in:  "",
    24  			err: false,
    25  			out: 0,
    26  		},
    27  		{
    28  			in:  "0",
    29  			err: false,
    30  			out: 0,
    31  		},
    32  		{
    33  			in:  "1b",
    34  			err: false,
    35  			out: 1,
    36  		},
    37  		{
    38  			in:  "100kb",
    39  			err: false,
    40  			out: 100 << 10,
    41  		},
    42  		{
    43  			in:  "100 KB",
    44  			err: false,
    45  			out: 100 << 10,
    46  		},
    47  		{
    48  			// ensure lowercase works
    49  			in:  "50mb",
    50  			err: false,
    51  			out: 50 << 20,
    52  		},
    53  		{
    54  			// ensure mixed capitalization works
    55  			in:  "50Mb",
    56  			err: false,
    57  			out: 50 << 20,
    58  		},
    59  		{
    60  			in:  "256GB",
    61  			err: false,
    62  			out: 256 << 30,
    63  		},
    64  	} {
    65  		t.Run(tc.in, func(t *testing.T) {
    66  			var bs ByteSize
    67  
    68  			err := bs.Set(tc.in)
    69  			if tc.err {
    70  				require.NotNil(t, err)
    71  			} else {
    72  				require.Nil(t, err)
    73  				require.Equal(t, tc.out, bs.Get().(int))
    74  			}
    75  
    76  		})
    77  	}
    78  }
    79  
    80  func Test_ByteSizeYAML(t *testing.T) {
    81  	for _, tc := range []struct {
    82  		in  string
    83  		err bool
    84  		out ByteSize
    85  	}{
    86  		{
    87  			in:  "256GB",
    88  			out: ByteSize(256 << 30),
    89  		},
    90  		{
    91  			in:  "abc",
    92  			err: true,
    93  		},
    94  	} {
    95  		t.Run(tc.in, func(t *testing.T) {
    96  			var out ByteSize
    97  			err := yaml.Unmarshal([]byte(tc.in), &out)
    98  			if tc.err {
    99  				require.NotNil(t, err)
   100  			} else {
   101  				require.Nil(t, err)
   102  				require.Equal(t, tc.out, out)
   103  			}
   104  		})
   105  	}
   106  }
   107  
   108  func Test_ByteSizeJSON(t *testing.T) {
   109  	for _, tc := range []struct {
   110  		in  string
   111  		err bool
   112  		out ByteSize
   113  	}{
   114  		{
   115  			in:  `{ "bytes": "256GB" }`,
   116  			out: ByteSize(256 << 30),
   117  		},
   118  		{
   119  			// JSON shouldn't allow to set integer as value for ByteSize field.
   120  			in:  `{ "bytes": 2.62144e+07 }`,
   121  			err: true,
   122  		},
   123  		{
   124  			in:  `{ "bytes": "abc" }`,
   125  			err: true,
   126  		},
   127  	} {
   128  		t.Run(tc.in, func(t *testing.T) {
   129  			var out struct {
   130  				Bytes ByteSize `json:"bytes"`
   131  			}
   132  			err := json.Unmarshal([]byte(tc.in), &out)
   133  			if tc.err {
   134  				require.NotNil(t, err)
   135  			} else {
   136  				require.Nil(t, err)
   137  				require.Equal(t, tc.out, out.Bytes)
   138  			}
   139  		})
   140  	}
   141  }