DomainWithUnit support existed query parameters

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

Hide file tree

Showing 3 changed files with 85 additions and 29 deletions.
diff --git a/consumer/option_test.go b/consumer/option_test.go
@@ -1,9 +1,7 @@
 package consumer
 
 import (
-	"fmt"
 	"reflect"
-	"strings"
 	"testing"
 )
 
@@ -15,7 +13,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 +22,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")
@@ -34,30 +32,58 @@ func TestWithNameServerDomain(t *testing.T) {
 }
 
 func TestWithNameServerDomainAndUnitName(t *testing.T) {
-	nameServerAddr := "http://127.0.0.1:8080/nameserver/addr"
 	unitName := "unsh"
-	suffix := fmt.Sprintf("-%s?nofix=1", unitName)
-
 	// test with two different orders
 	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
-		opt := defaultPullConsumerOptions()
-		WithNameServerDomain(nameServerAddr)(&opt)
+		addr := "http://127.0.0.1:8080/nameserver/addr"
+		opt := defaultPushConsumerOptions()
+		WithNameServerDomain(addr)(&opt)
+		WithUnitName(unitName)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1"
+		if domainStr != expectedAddr {
+			t.Errorf("consumer option WithNameServerDomain & WithUnitName. want:%s, got=%s", expectedAddr, domainStr)
+		}
+	})
+
+	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr"
+		opt := defaultPushConsumerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(addr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1"
+		if domainStr != expectedAddr {
+			t.Errorf("consumer option WithUnitName & WithNameServerDomain. want:%s, got=%s", expectedAddr, domainStr)
+		}
+	})
+
+	// test with two different orders - name server with query string
+	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		opt := defaultPushConsumerOptions()
+		WithNameServerDomain(addr)(&opt)
 		WithUnitName(unitName)(&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)
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1&labels=abc"
+		if domainStr != expectedAddr {
+			t.Errorf("consumer option WithNameServerDomain & WithUnitName. want:%s, got=%s", expectedAddr, domainStr)
 		}
 	})
 
 	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
-		opt := defaultPullConsumerOptions()
-		WithNameServerDomain(nameServerAddr)(&opt)
+		addr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
+		opt := defaultPushConsumerOptions()
 		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(addr)(&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)
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1&labels=abc"
+		if domainStr != expectedAddr {
+			t.Errorf("consumer option WithUnitName & WithNameServerDomain. want:%s, got=%s", expectedAddr, domainStr)
 		}
 	})
 }
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
@@ -1,9 +1,7 @@
 package producer
 
 import (
-	"fmt"
 	"reflect"
-	"strings"
 	"testing"
 )
 
@@ -19,7 +17,7 @@ func TestWithUnitName(t *testing.T) {
 	unitName := "unsh"
 	WithUnitName(unitName)(&opt)
 	if opt.UnitName != unitName {
-		t.Errorf("consumer option WithUnitName. want:%s, got=%s", unitName, opt.UnitName)
+		t.Errorf("producer option WithUnitName. want:%s, got=%s", unitName, opt.UnitName)
 	}
 }
 
@@ -29,35 +27,63 @@ func TestWithNameServerDomain(t *testing.T) {
 	WithNameServerDomain(nameServerAddr)(&opt)
 	domainStr := getFieldString(opt.Resolver, "domain")
 	if domainStr != nameServerAddr {
-		t.Errorf("consumer option WithUnitName. want:%s, got=%s", nameServerAddr, domainStr)
+		t.Errorf("producer option WithUnitName. want:%s, got=%s", nameServerAddr, domainStr)
 	}
 }
 
 func TestWithNameServerDomainAndUnitName(t *testing.T) {
-	nameServerAddr := "http://127.0.0.1:8080/nameserver/addr"
 	unitName := "unsh"
-	suffix := fmt.Sprintf("-%s?nofix=1", unitName)
-
 	// test with two different orders
 	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr"
+		opt := defaultProducerOptions()
+		WithNameServerDomain(addr)(&opt)
+		WithUnitName(unitName)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1"
+		if domainStr != expectedAddr {
+			t.Errorf("producer option WithNameServerDomain & WithUnitName. want:%s, got=%s", expectedAddr, domainStr)
+		}
+	})
+
+	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr"
+		opt := defaultProducerOptions()
+		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(addr)(&opt)
+
+		domainStr := getFieldString(opt.Resolver, "domain")
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1"
+		if domainStr != expectedAddr {
+			t.Errorf("producer option WithUnitName & WithNameServerDomain. want:%s, got=%s", expectedAddr, domainStr)
+		}
+	})
+
+	// test with two different orders - name server with query string
+	t.Run("WithNameServerDomain & WithUnitName", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
 		opt := defaultProducerOptions()
-		WithNameServerDomain(nameServerAddr)(&opt)
+		WithNameServerDomain(addr)(&opt)
 		WithUnitName(unitName)(&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)
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1&labels=abc"
+		if domainStr != expectedAddr {
+			t.Errorf("producer option WithNameServerDomain & WithUnitName. want:%s, got=%s", expectedAddr, domainStr)
 		}
 	})
 
 	t.Run("WithUnitName & WithNameServerDomain", func(t *testing.T) {
+		addr := "http://127.0.0.1:8080/nameserver/addr?labels=abc"
 		opt := defaultProducerOptions()
-		WithNameServerDomain(nameServerAddr)(&opt)
 		WithUnitName(unitName)(&opt)
+		WithNameServerDomain(addr)(&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)
+		expectedAddr := "http://127.0.0.1:8080/nameserver/addr-unsh?nofix=1&labels=abc"
+		if domainStr != expectedAddr {
+			t.Errorf("producer option WithUnitName & WithNameServerDomain. want:%s, got=%s", expectedAddr, domainStr)
 		}
 	})
 }