Improve version switching for Travis builds

client/remote_test.go

@@ -12,13 +12,18 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func getTestConfig(writer io.Writer) *Config {
-	config := &Config{Writer: writer}
+func getVersion() string {
 	travis := os.Getenv("TRAVIS")
 	if travis != "" {
-		config.Version = "travis"
-	} else {
-		config.Version = "canary"
+		return "travis"
+	}
+	return "canary"
+}
+
+func getTestConfig(writer io.Writer) *Config {
+	config := &Config{
+		Writer:  writer,
+		Version: getVersion(),
 	}
 	return config
 }
@@ -94,13 +99,13 @@ func TestKeyGen(t *testing.T) {
 	remote := getInstrumentedTestRemote()
 	script, err := ioutil.ReadFile("bootstrap/token.sh")
 	assert.Nil(t, err)
-	tokenScript := fmt.Sprintf(string(script), "canary")
+	tokenScript := fmt.Sprintf(string(script), getVersion())
 
 	// Make sure the right command is run.
 	session := mockSSHRunner{r: remote}
 
 	// Make sure the right command is run.
-	_, err = remote.GetDaemonAPIToken(&session, "canary")
+	_, err = remote.GetDaemonAPIToken(&session, getVersion())
 	assert.Nil(t, err)
 	assert.Equal(t, session.Calls[0], tokenScript)
 }
@@ -112,14 +117,14 @@ func TestBootstrap(t *testing.T) {
 
 	script, err := ioutil.ReadFile("bootstrap/daemon-up.sh")
 	assert.Nil(t, err)
-	daemonScript := fmt.Sprintf(string(script), "canary", "8081")
+	daemonScript := fmt.Sprintf(string(script), getVersion(), "8081")
 
 	keyScript, err := ioutil.ReadFile("bootstrap/keygen.sh")
 	assert.Nil(t, err)
 
 	script, err = ioutil.ReadFile("bootstrap/token.sh")
 	assert.Nil(t, err)
-	tokenScript := fmt.Sprintf(string(script), "canary")
+	tokenScript := fmt.Sprintf(string(script), getVersion())
 
 	var writer bytes.Buffer
 	session := mockSSHRunner{r: remote}