pkg/flags: fix UniqueURLs'Set to remove duplicates in UniqueURLs'uss

From the name of func 'UniqueURLsFromFlag', we can tell that UniqueURLs'uss should not have duplicates. The current implemention of UniqueURLs'Set has a bug to make it unique. Fixes: etcd-io#16307. Signed-off-by: Jes Cok <[email protected]>
ahrtr · Jul 26, 2023 · 5df43d9 · 5df43d9
1 parent 7ba4add
commit 5df43d9
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/pkg/flags/unique_urls.go b/pkg/flags/unique_urls.go
@@ -50,7 +50,11 @@ func (us *UniqueURLs) Set(s string) error {
 	us.Values = make(map[string]struct{})
 	us.uss = make([]url.URL, 0)
 	for _, v := range ss {
-		us.Values[v.String()] = struct{}{}
+		x := v.String()
+		if _, exists := us.Values[x]; exists {
+			continue
+		}
+		us.Values[x] = struct{}{}
 		us.uss = append(us.uss, v)
 	}
 	return nil

diff --git a/pkg/flags/unique_urls_test.go b/pkg/flags/unique_urls_test.go
@@ -15,8 +15,11 @@
 package flags
 
 import (
-	"reflect"
+	"flag"
+	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewUniqueURLsWithExceptions(t *testing.T) {
@@ -83,11 +86,29 @@ func TestNewUniqueURLsWithExceptions(t *testing.T) {
 	}
 	for i := range tests {
 		uv := NewUniqueURLsWithExceptions(tests[i].s, tests[i].exception)
-		if !reflect.DeepEqual(tests[i].exp, uv.Values) {
-			t.Fatalf("#%d: expected %+v, got %+v", i, tests[i].exp, uv.Values)
-		}
-		if uv.String() != tests[i].rs {
-			t.Fatalf("#%d: expected %q, got %q", i, tests[i].rs, uv.String())
-		}
+		require.Equal(t, tests[i].exp, uv.Values)
+		require.Equal(t, uv.String(), tests[i].rs)
+	}
+}
+
+func TestUniqueURLsFromFlag(t *testing.T) {
+	const name = "test"
+	urls := []string{
+		"https://1.2.3.4:1",
+		"https://1.2.3.4:2",
+		"https://1.2.3.4:3",
+		"https://1.2.3.4:1",
+	}
+	fs := flag.NewFlagSet(name, flag.ExitOnError)
+	u := NewUniqueURLsWithExceptions(strings.Join(urls, ","))
+	fs.Var(u, name, "usage")
+	uss := UniqueURLsFromFlag(fs, name)
+
+	require.Equal(t, len(u.Values), len(uss))
+
+	um := make(map[string]struct{})
+	for _, x := range uss {
+		um[x.String()] = struct{}{}
 	}
+	require.Equal(t, u.Values, um)
 }