Browse Source

fix from comments

Barak Michener 11 years ago
parent
commit
04d9f848a7
2 changed files with 23 additions and 23 deletions
  1. 21 22
      etcdmain/etcd.go
  2. 2 1
      etcdmain/etcd_test.go

+ 21 - 22
etcdmain/etcd.go

@@ -453,33 +453,32 @@ func genDNSClusterString(defaultToken string) (string, string, error) {
 	targetName := make(map[string]int)
 	stringParts := make([]string, 0)
 	tempName := int(0)
-	_, addrs, err := lookupSRV("etcd-server-ssl", "tcp", *dnsCluster)
-	if err != nil {
-		return "", "", err
-	}
-	for _, srv := range addrs {
-		var v int
-		var ok bool
-		if v, ok = targetName[srv.Target]; !ok {
-			v = tempName
-			targetName[srv.Target] = v
-			tempName += 1
+
+	updateNodeMap := func(service, prefix string) error {
+		_, addrs, err := lookupSRV(service, "tcp", *dnsCluster)
+		if err != nil {
+			return err
 		}
-		stringParts = append(stringParts, fmt.Sprintf("%d=https://%s:%d", v, srv.Target, srv.Port))
+		for _, srv := range addrs {
+			var v int
+			var ok bool
+			if v, ok = targetName[srv.Target]; !ok {
+				v = tempName
+				targetName[srv.Target] = v
+				tempName += 1
+			}
+			stringParts = append(stringParts, fmt.Sprintf("%d=%s%s:%d", v, prefix, srv.Target, srv.Port))
+		}
+		return nil
 	}
-	_, addrs, err = lookupSRV("etcd-server", "tcp", *dnsCluster)
+
+	err := updateNodeMap("etcd-server-ssl", "https://")
 	if err != nil {
 		return "", "", err
 	}
-	for _, srv := range addrs {
-		var v int
-		var ok bool
-		if v, ok = targetName[srv.Target]; !ok {
-			v = tempName
-			targetName[srv.Target] = v
-			tempName += 1
-		}
-		stringParts = append(stringParts, fmt.Sprintf("%d=http://%s:%d", v, srv.Target, srv.Port))
+	err = updateNodeMap("etcd-server", "http://")
+	if err != nil {
+		return "", "", err
 	}
 	return strings.Join(stringParts, ","), defaultToken, nil
 }

+ 2 - 1
etcdmain/etcd_test.go

@@ -100,12 +100,13 @@ func TestGenDNSClusterString(t *testing.T) {
 			}
 			return "", nil, errors.New("Unkown service in mock")
 		}
+		defer func() { lookupSRV = net.LookupSRV }()
 		str, token, err := genDNSClusterString("token")
 		if err != nil {
 			t.Fatalf("%d: err: %#v", i, err)
 		}
 		if token != "token" {
-			t.Error("Token doesn't match default token")
+			t.Errorf("%d: token: %s", i, token)
 		}
 		if str != tt.expected {
 			t.Errorf("#%d: cluster = %s, want %s", i, str, tt.expected)