From 9831d712b16dcfbb62d376aa45222fcde12f3dca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fan=20Can=20Bak=C4=B1r?= Date: Fri, 21 Jul 2023 07:13:03 +0000 Subject: [PATCH] use StringSlice --- ratelimit_var.go | 19 +++++++++++-------- ratelimit_var_test.go | 22 +++++++++++++++++++--- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/ratelimit_var.go b/ratelimit_var.go index 8cb57e6..0c310ed 100644 --- a/ratelimit_var.go +++ b/ratelimit_var.go @@ -69,28 +69,31 @@ func (rateLimitMap RateLimitMap) String() string { var items string for k, v := range rateLimitMap.kv { - items += fmt.Sprintf("\"%s\"=\"%s\"%s", k, v.Duration.String(), kvSep) + items += fmt.Sprintf("\"%s\":\"%d/%s\",", k, v.MaxCount, v.Duration.String()) } - defaultBuilder.WriteString(stringsutil.TrimSuffixAny(items, ",", "=")) + defaultBuilder.WriteString(stringsutil.TrimSuffixAny(items, ",", ":")) defaultBuilder.WriteString("}") return defaultBuilder.String() } // RateLimitMapVar adds a ratelimit flag with a longname -func (flagSet *FlagSet) RateLimitMapVar(field *RateLimitMap, long string, defaultValue []string, usage string) *FlagData { - return flagSet.RateLimitMapVarP(field, long, "", defaultValue, usage) +func (flagSet *FlagSet) RateLimitMapVar(field *RateLimitMap, long string, defaultValue []string, usage string, options Options) *FlagData { + return flagSet.RateLimitMapVarP(field, long, "", defaultValue, usage, options) } // RateLimitMapVarP adds a ratelimit flag with a short name and long name. // It is equivalent to RateLimitMapVar, and also allows specifying ratelimits in days (e.g., "hackertarget=2/d" 2 requests per day, which is equivalent to 24h). -func (flagSet *FlagSet) RateLimitMapVarP(field *RateLimitMap, long, short string, defaultValue []string, usage string) *FlagData { +func (flagSet *FlagSet) RateLimitMapVarP(field *RateLimitMap, long, short string, defaultValue StringSlice, usage string, options Options) *FlagData { if field == nil { panic(fmt.Errorf("field cannot be nil for flag -%v", long)) } - for _, item := range defaultValue { - if err := field.Set(item); err != nil { - panic(fmt.Errorf("failed to set default value for flag -%v: %v", long, err)) + for _, defaultItem := range defaultValue { + values, _ := ToStringSlice(defaultItem, options) + for _, value := range values { + if err := field.Set(value); err != nil { + panic(fmt.Errorf("failed to set default value for flag -%v: %v", long, err)) + } } } diff --git a/ratelimit_var_test.go b/ratelimit_var_test.go index 2ef9c2d..a68913e 100644 --- a/ratelimit_var_test.go +++ b/ratelimit_var_test.go @@ -14,7 +14,7 @@ func TestRateLimitMapVar(t *testing.T) { var rateLimitMap RateLimitMap flagSet := NewFlagSet() flagSet.CreateGroup("Config", "Config", - flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/ms"}, "rate limits"), + flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/ms"}, "rate limits", CommaSeparatedStringSliceOptions), ) os.Args = []string{ os.Args[0], @@ -25,11 +25,27 @@ func TestRateLimitMapVar(t *testing.T) { tearDown(t.Name()) }) + t.Run("multiple-default-value", func(t *testing.T) { + var rateLimitMap RateLimitMap + flagSet := NewFlagSet() + flagSet.CreateGroup("Config", "Config", + flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", []string{"hackertarget=1/s,github=1/ms"}, "rate limits", CommaSeparatedStringSliceOptions), + ) + os.Args = []string{ + os.Args[0], + } + err := flagSet.Parse() + assert.Nil(t, err) + assert.Equal(t, RateLimit{MaxCount: 1, Duration: time.Second}, rateLimitMap.AsMap()["hackertarget"]) + assert.Equal(t, RateLimit{MaxCount: 1, Duration: time.Millisecond}, rateLimitMap.AsMap()["github"]) + tearDown(t.Name()) + }) + t.Run("valid-rate-limit", func(t *testing.T) { var rateLimitMap RateLimitMap flagSet := NewFlagSet() flagSet.CreateGroup("Config", "Config", - flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits"), + flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits", CommaSeparatedStringSliceOptions), ) os.Args = []string{ os.Args[0], @@ -45,7 +61,7 @@ func TestRateLimitMapVar(t *testing.T) { var rateLimitMap RateLimitMap flagSet := NewFlagSet() flagSet.CreateGroup("Config", "Config", - flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits"), + flagSet.RateLimitMapVarP(&rateLimitMap, "rate-limits", "rls", nil, "rate limits", CommaSeparatedStringSliceOptions), ) os.Args = []string{ os.Args[0],