allow custom url on token

Signed-off-by: Victor Vieux <vieux@docker.com>

discovery/token/token.go

@@ -14,6 +14,7 @@ import (
 const DISCOVERY_URL = "https://discovery-stage.hub.docker.com/v1"
 
 type TokenDiscoveryService struct {
+	url   string
 	token string
 }
 
@@ -21,13 +22,23 @@ func init() {
 	discovery.Register("token", Init)
 }
 
-func Init(token string) (discovery.DiscoveryService, error) {
-	return TokenDiscoveryService{token: token}, nil
+func Init(urltoken string) (discovery.DiscoveryService, error) {
+	if i := strings.LastIndex(urltoken, "/"); i != -1 {
+		return TokenDiscoveryService{url: "https://" + urltoken[:i], token: urltoken[i+1:]}, nil
+	}
+
+	return TokenDiscoveryService{url: DISCOVERY_URL, token: urltoken}, nil
+}
+func New(url string) *TokenDiscoveryService {
+	if url != "" {
+		return &TokenDiscoveryService{url: url}
+	}
+	return &TokenDiscoveryService{url: DISCOVERY_URL}
 }
 
 // FetchNodes returns the node for the discovery service at the specified endpoint
 func (s TokenDiscoveryService) Fetch() ([]string, error) {
-	resp, err := http.Get(fmt.Sprintf("%s/%s/%s", DISCOVERY_URL, "clusters", s.token))
+	resp, err := http.Get(fmt.Sprintf("%s/%s/%s", s.url, "clusters", s.token))
 	if err != nil {
 		return nil, err
 	}
@@ -54,14 +65,14 @@ func (s TokenDiscoveryService) Watch(heartbeat int) <-chan time.Time {
 func (s TokenDiscoveryService) Register(addr string) error {
 	buf := strings.NewReader(addr)
 
-	_, err := http.Post(fmt.Sprintf("%s/%s/%s", DISCOVERY_URL,
+	_, err := http.Post(fmt.Sprintf("%s/%s/%s", s.url,
 		"clusters", s.token), "application/json", buf)
 	return err
 }
 
 // CreateCluster returns a unique cluster token
-func CreateCluster() (string, error) {
-	resp, err := http.Post(fmt.Sprintf("%s/%s", DISCOVERY_URL, "clusters"), "", nil)
+func (s TokenDiscoveryService) CreateCluster() (string, error) {
+	resp, err := http.Post(fmt.Sprintf("%s/%s", s.url, "clusters"), "", nil)
 	if err != nil {
 		return "", err
 	}

discovery/token/token_test.go

@@ -1,28 +1,34 @@
 package token
 
-import "testing"
+import (
+	"testing"
 
-func TestRegister(t *testing.T) {
-	discovery := TokenDiscoveryService{token: "TEST_TOKEN"}
-	expected := "127.0.0.1:2675"
-	if err := discovery.RegisterNode(expected); err != nil {
-		t.Fatal(err)
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInit(t *testing.T) {
+	discovery, _ := Init("token")
+	if dtoken, ok := discovery.(TokenDiscoveryService); ok {
+		assert.Equal(t, dtoken.token, "token")
+		assert.Equal(t, dtoken.url, DISCOVERY_URL)
 	}
 
-	addrs, err := discovery.FetchNodes()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if len(addrs) != 1 {
-		t.Fatalf("expected addr len == 1, got len = %d", len(addrs))
-	}
-
-	if addrs[0] != expected {
-		t.Fatalf("expected addr %q but received %q", expected, addrs[0])
-	}
-
-	if err = discovery.RegisterNode(expected); err != nil {
-		t.Fatal(err)
+	discovery, _ = Init("custom/path/token")
+	if dtoken, ok := discovery.(TokenDiscoveryService); ok {
+		assert.Equal(t, dtoken.token, "token")
+		assert.Equal(t, dtoken.url, "https://custom/path")
 	}
 }
+
+func TestRegister(t *testing.T) {
+	discovery := TokenDiscoveryService{token: "TEST_TOKEN", url: DISCOVERY_URL}
+	expected := "127.0.0.1:2675"
+	assert.NoError(t, discovery.Register(expected))
+
+	addrs, err := discovery.Fetch()
+	assert.NoError(t, err)
+	assert.Equal(t, len(addrs), 1)
+	assert.Equal(t, addrs[0], expected)
+
+	assert.NoError(t, discovery.Register(expected))
+}

main.go

@@ -84,7 +84,7 @@ func main() {
 			ShortName: "c",
 			Usage:     "create a cluster",
 			Action: func(c *cli.Context) {
-				token, err := token.CreateCluster()
+				token, err := token.New("").CreateCluster()
 				if err != nil {
 					log.Fatal(err)
 				}