diff --git a/etcdmain/etcd.go b/etcdmain/etcd.go index a56615a82..c9810bf85 100644 --- a/etcdmain/etcd.go +++ b/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) + + updateNodeMap := func(service, prefix string) error { + _, addrs, err := lookupSRV(service, "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 + } + stringParts = append(stringParts, fmt.Sprintf("%d=%s%s:%d", v, prefix, srv.Target, srv.Port)) + } + return nil + } + + 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=https://%s:%d", v, srv.Target, srv.Port)) - } - _, addrs, err = lookupSRV("etcd-server", "tcp", *dnsCluster) + err = updateNodeMap("etcd-server", "http://") 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)) - } return strings.Join(stringParts, ","), defaultToken, nil } diff --git a/etcdmain/etcd_test.go b/etcdmain/etcd_test.go index ac461cb2f..517cfbb62 100644 --- a/etcdmain/etcd_test.go +++ b/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)