Add a simple test and mock for genDNS

etcdmain/etcd.go

@@ -103,6 +103,8 @@ var (
 		"v",
 		"vv",
 	}
+
+	lookupSRV = net.LookupSRV
 )
 
 func init() {
@@ -422,7 +424,7 @@ func setupCluster(apurls []url.URL) (*etcdserver.Cluster, error) {
 		clusterStr := genClusterString(*name, apurls)
 		cls, err = etcdserver.NewClusterFromString(*durl, clusterStr)
 	case set["dns-cluster-domain"]:
-		clusterStr, clusterToken, err := genDnsClusterString(*initialClusterToken)
+		clusterStr, clusterToken, err := genDNSClusterString(*initialClusterToken)
 		if err != nil {
 			return nil, err
 		}
@@ -447,11 +449,11 @@ func genClusterString(name string, urls types.URLs) string {
 // TODO(barakmich): Currently ignores priority and weight (as they don't make as much sense for a bootstrap)
 // Also doesn't do any lookups for the token (though it could)
 // Also sees hostnames and IPs as separate -- use one or the other for consistency.
-func genDnsClusterString(defaultToken string) (string, string, error) {
+func genDNSClusterString(defaultToken string) (string, string, error) {
 	targetName := make(map[string]int)
 	stringParts := make([]string, 0)
 	tempName := int(0)
-	_, addrs, err := net.LookupSRV("etcd-server-ssl", "tcp", *dnsCluster)
+	_, addrs, err := lookupSRV("etcd-server-ssl", "tcp", *dnsCluster)
 	if err != nil {
 		return "", "", err
 	}
@@ -465,7 +467,7 @@ func genDnsClusterString(defaultToken string) (string, string, error) {
 		}
 		stringParts = append(stringParts, fmt.Sprintf("%d=https://%s:%d", v, srv.Target, srv.Port))
 	}
-	_, addrs, err = net.LookupSRV("etcd-server", "tcp", *dnsCluster)
+	_, addrs, err = lookupSRV("etcd-server", "tcp", *dnsCluster)
 	if err != nil {
 		return "", "", err
 	}
@@ -479,5 +481,5 @@ func genDnsClusterString(defaultToken string) (string, string, error) {
 		}
 		stringParts = append(stringParts, fmt.Sprintf("%d=http://%s:%d", v, srv.Target, srv.Port))
 	}
-	return defaultToken, strings.Join(stringParts, ","), nil
+	return strings.Join(stringParts, ","), defaultToken, nil
 }

etcdmain/etcd_test.go

@@ -17,6 +17,8 @@
 package etcdmain
 
 import (
+	"errors"
+	"net"
 	"net/url"
 	"testing"
 
@@ -54,3 +56,59 @@ func TestGenClusterString(t *testing.T) {
 		}
 	}
 }
+
+func TestGenDNSClusterString(t *testing.T) {
+	tests := []struct {
+		withSSL    []*net.SRV
+		withoutSSL []*net.SRV
+		expected   string
+	}{
+		{
+			[]*net.SRV{},
+			[]*net.SRV{},
+			"",
+		},
+		{
+			[]*net.SRV{
+				&net.SRV{Target: "10.0.0.1", Port: 2480},
+				&net.SRV{Target: "10.0.0.2", Port: 2480},
+				&net.SRV{Target: "10.0.0.3", Port: 2480},
+			},
+			[]*net.SRV{},
+			"0=https://10.0.0.1:2480,1=https://10.0.0.2:2480,2=https://10.0.0.3:2480",
+		},
+		{
+			[]*net.SRV{
+				&net.SRV{Target: "10.0.0.1", Port: 2480},
+				&net.SRV{Target: "10.0.0.2", Port: 2480},
+				&net.SRV{Target: "10.0.0.3", Port: 2480},
+			},
+			[]*net.SRV{
+				&net.SRV{Target: "10.0.0.1", Port: 7001},
+			},
+			"0=https://10.0.0.1:2480,1=https://10.0.0.2:2480,2=https://10.0.0.3:2480,0=http://10.0.0.1:7001",
+		},
+	}
+
+	for i, tt := range tests {
+		lookupSRV = func(service string, proto string, domain string) (string, []*net.SRV, error) {
+			if service == "etcd-server-ssl" {
+				return "", tt.withSSL, nil
+			}
+			if service == "etcd-server" {
+				return "", tt.withoutSSL, nil
+			}
+			return "", nil, errors.New("Unkown service in mock")
+		}
+		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")
+		}
+		if str != tt.expected {
+			t.Errorf("#%d: cluster = %s, want %s", i, str, tt.expected)
+		}
+	}
+}