DomainWithUnit support existed query parameters

apache · May 18, 2023 · 0071322 · 0071322
1 parent 8c4e388
commit 0071322
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 5 deletions.
diff --git a/consumer/option_test.go b/consumer/option_test.go
@@ -15,7 +15,7 @@ func getFieldString(obj interface{}, field string) string {
 }
 
 func TestWithUnitName(t *testing.T) {
-	opt := defaultPullConsumerOptions()
+	opt := defaultPushConsumerOptions()
 	unitName := "unsh"
 	WithUnitName(unitName)(&opt)
 	if opt.UnitName != unitName {
@@ -24,7 +24,7 @@ func TestWithUnitName(t *testing.T) {
 }
 
 func TestWithNameServerDomain(t *testing.T) {
-	opt := defaultPullConsumerOptions()
+	opt := defaultPushConsumerOptions()
 	nameServerAddr := "http://127.0.0.1:8080/nameserver/addr"
 	WithNameServerDomain(nameServerAddr)(&opt)
 	domainStr := getFieldString(opt.Resolver, "domain")
@@ -40,7 +40,7 @@ func TestWithNameServerDomainAndUnitName(t *testing.T) {
 
 	// test with two different orders
 	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
-		opt := defaultPullConsumerOptions()
+		opt := defaultPushConsumerOptions()
 		WithNameServerDomain(nameServerAddr)(&opt)
 		WithUnitName(unitName)(&opt)
 
@@ -51,7 +51,22 @@ func TestWithNameServerDomainAndUnitName(t *testing.T) {
 	})
 
 	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
-		opt := defaultPullConsumerOptions()
+		opt := defaultPushConsumerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(nameServerAddr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		if !strings.Contains(domainStr, nameServerAddr) || !strings.Contains(domainStr, suffix) {
+			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
+		}
+	})
+
+	// test with two different orders - name server with query string
+	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
+		nameServerAddr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		suffix := fmt.Sprintf("-%s?nofix=1&labels=abc", unitName)
+
+		opt := defaultPushConsumerOptions()
 		WithNameServerDomain(nameServerAddr)(&opt)
 		WithUnitName(unitName)(&opt)
 
@@ -60,4 +75,18 @@ func TestWithNameServerDomainAndUnitName(t *testing.T) {
 			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
 		}
 	})
+
+	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		nameServerAddr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		suffix := fmt.Sprintf("-%s?nofix=1&labels=abc", unitName)
+
+		opt := defaultPushConsumerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(nameServerAddr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		if !strings.Contains(domainStr, nameServerAddr) || !strings.Contains(domainStr, suffix) {
+			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
+		}
+	})
 }
diff --git a/primitive/nsresolver.go b/primitive/nsresolver.go
@@ -118,7 +118,11 @@ func (h *HttpResolver) DomainWithUnit(unitName string) {
 	if strings.Contains(h.domain, "?nofix=1") {
 		return
 	}
-	h.domain = fmt.Sprintf("%s-%s?nofix=1", h.domain, unitName)
+	if strings.Contains(h.domain, "?") {
+		h.domain = strings.Replace(h.domain, "?", fmt.Sprintf("-%s?nofix=1&", unitName), 1)
+	} else {
+		h.domain = fmt.Sprintf("%s-%s?nofix=1", h.domain, unitName)
+	}
 }
 
 func (h *HttpResolver) Resolve() []string {

diff --git a/producer/option_test.go b/producer/option_test.go
@@ -51,6 +51,21 @@ func TestWithNameServerDomainAndUnitName(t *testing.T) {
 	})
 
 	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		opt := defaultProducerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(nameServerAddr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		if !strings.Contains(domainStr, nameServerAddr) || !strings.Contains(domainStr, suffix) {
+			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
+		}
+	})
+
+	// test with two different orders - name server with query string
+	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
+		nameServerAddr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		suffix := fmt.Sprintf("-%s?nofix=1&labels=abc", unitName)
+
 		opt := defaultProducerOptions()
 		WithNameServerDomain(nameServerAddr)(&opt)
 		WithUnitName(unitName)(&opt)
@@ -60,4 +75,18 @@ func TestWithNameServerDomainAndUnitName(t *testing.T) {
 			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
 		}
 	})
+
+	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		nameServerAddr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		suffix := fmt.Sprintf("-%s?nofix=1&labels=abc", unitName)
+
+		opt := defaultProducerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(nameServerAddr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		if !strings.Contains(domainStr, nameServerAddr) || !strings.Contains(domainStr, suffix) {
+			t.Errorf("consumer option should contains %s and %s", nameServerAddr, suffix)
+		}
+	})
 }