github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/util/validation/notifications_limit_flag_test.go (about)

     1  package validation
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  	"github.com/stretchr/testify/require"
    10  	"gopkg.in/yaml.v2"
    11  )
    12  
    13  func TestNotificationLimitsMap(t *testing.T) {
    14  	for name, tc := range map[string]struct {
    15  		args     []string
    16  		expected NotificationRateLimitMap
    17  		error    string
    18  	}{
    19  		"basic test": {
    20  			args: []string{"-map-flag", "{\"email\": 100 }"},
    21  			expected: NotificationRateLimitMap{
    22  				"email": 100,
    23  			},
    24  		},
    25  
    26  		"unknown integration": {
    27  			args:  []string{"-map-flag", "{\"unknown\": 200 }"},
    28  			error: "invalid value \"{\\\"unknown\\\": 200 }\" for flag -map-flag: unknown integration name: unknown",
    29  		},
    30  
    31  		"parsing error": {
    32  			args:  []string{"-map-flag", "{\"hello\": ..."},
    33  			error: "invalid value \"{\\\"hello\\\": ...\" for flag -map-flag: invalid character '.' looking for beginning of value",
    34  		},
    35  	} {
    36  		t.Run(name, func(t *testing.T) {
    37  			v := NotificationRateLimitMap{}
    38  
    39  			fs := flag.NewFlagSet("test", flag.ContinueOnError)
    40  			fs.SetOutput(&bytes.Buffer{}) // otherwise errors would go to stderr.
    41  			fs.Var(v, "map-flag", "Map flag, you can pass JSON into this")
    42  			err := fs.Parse(tc.args)
    43  
    44  			if tc.error != "" {
    45  				require.NotNil(t, err)
    46  				assert.Equal(t, tc.error, err.Error())
    47  			} else {
    48  				assert.NoError(t, err)
    49  				assert.Equal(t, tc.expected, v)
    50  			}
    51  		})
    52  	}
    53  }
    54  
    55  type TestStruct struct {
    56  	Flag NotificationRateLimitMap `yaml:"flag"`
    57  }
    58  
    59  func TestNotificationsLimitMapYaml(t *testing.T) {
    60  
    61  	var testStruct TestStruct
    62  	testStruct.Flag = map[string]float64{}
    63  
    64  	require.NoError(t, testStruct.Flag.Set("{\"email\": 500 }"))
    65  	expected := []byte(`flag:
    66    email: 500
    67  `)
    68  
    69  	actual, err := yaml.Marshal(testStruct)
    70  	require.NoError(t, err)
    71  	assert.Equal(t, expected, actual)
    72  
    73  	var actualStruct TestStruct
    74  	actualStruct.Flag = NotificationRateLimitMap{} // must be set, otherwise unmarshalling panics.
    75  
    76  	err = yaml.Unmarshal(expected, &actualStruct)
    77  	require.NoError(t, err)
    78  	assert.Equal(t, testStruct, actualStruct)
    79  }
    80  
    81  func TestUnknownIntegrationWhenLoadingYaml(t *testing.T) {
    82  	var s TestStruct
    83  	s.Flag = NotificationRateLimitMap{} // must be set, otherwise unmarshalling panics.
    84  
    85  	yamlInput := `flag:
    86    unknown_integration: 500
    87  `
    88  
    89  	err := yaml.Unmarshal([]byte(yamlInput), &s)
    90  	require.NotNil(t, err)
    91  	require.Equal(t, "unknown integration name: unknown_integration", err.Error())
    92  }
    93  
    94  func TestWrongYamlStructureWhenLoadingYaml(t *testing.T) {
    95  	var s TestStruct
    96  	s.Flag = NotificationRateLimitMap{} // must be set, otherwise unmarshalling panics.
    97  
    98  	yamlInput := `flag:
    99    email:
   100      rate_limit: 7777
   101      burst_size: 7777
   102  `
   103  
   104  	err := yaml.Unmarshal([]byte(yamlInput), &s)
   105  	require.NotNil(t, err)
   106  	require.Contains(t, err.Error(), "cannot unmarshal !!map into float64")
   107  }