From a09f90f968f2850ffbcb6386aa6a5e512efa6048 Mon Sep 17 00:00:00 2001
From: "Jonathan Leibiusky @xetorthio" <ionathan@gmail.com>
Date: Mon, 5 Jun 2017 10:08:03 -0300
Subject: [PATCH] Allow custom hostnames when creating instances

---
 pwd/instance.go      | 37 +++++++++++++++++++++-------------
 pwd/instance_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/pwd/instance.go b/pwd/instance.go
index a9ec36a..5f3ed81 100644
--- a/pwd/instance.go
+++ b/pwd/instance.go
@@ -60,6 +60,7 @@ type Instance struct {
 type InstanceConfig struct {
 	ImageName  string
 	Alias      string
+	Hostname   string
 	ServerCert []byte
 	ServerKey  []byte
 	CACert     []byte
@@ -178,35 +179,43 @@ func (p *pwd) InstanceDelete(session *Session, instance *Instance) error {
 	return nil
 }
 
+func (p *pwd) checkHostnameExists(session *Session, hostname string) bool {
+	containerName := fmt.Sprintf("%s_%s", session.Id[:8], hostname)
+	exists := false
+	for _, instance := range session.Instances {
+		if instance.Name == containerName {
+			exists = true
+			break
+		}
+	}
+	return exists
+}
+
 func (p *pwd) InstanceNew(session *Session, conf InstanceConfig) (*Instance, error) {
 	if conf.ImageName == "" {
 		conf.ImageName = config.GetDindImageName()
 	}
 	log.Printf("NewInstance - using image: [%s]\n", conf.ImageName)
 
-	var nodeName string
-	var containerName string
-	for i := 1; ; i++ {
-		nodeName = fmt.Sprintf("node%d", i)
-		containerName = fmt.Sprintf("%s_%s", session.Id[:8], nodeName)
-		exists := false
-		for _, instance := range session.Instances {
-			if instance.Name == containerName {
-				exists = true
+	if conf.Hostname == "" {
+		var nodeName string
+		for i := 1; ; i++ {
+			nodeName = fmt.Sprintf("node%d", i)
+			exists := p.checkHostnameExists(session, nodeName)
+			if !exists {
 				break
 			}
 		}
-		if !exists {
-			break
-		}
+		conf.Hostname = nodeName
 	}
+	containerName := fmt.Sprintf("%s_%s", session.Id[:8], conf.Hostname)
 
 	opts := docker.CreateContainerOpts{
 		Image:         conf.ImageName,
 		SessionId:     session.Id,
 		PwdIpAddress:  session.PwdIpAddress,
 		ContainerName: containerName,
-		Hostname:      nodeName,
+		Hostname:      conf.Hostname,
 		ServerCert:    conf.ServerCert,
 		ServerKey:     conf.ServerKey,
 		CACert:        conf.CACert,
@@ -229,7 +238,7 @@ func (p *pwd) InstanceNew(session *Session, conf InstanceConfig) (*Instance, err
 	instance.Image = opts.Image
 	instance.IP = ip
 	instance.Name = containerName
-	instance.Hostname = nodeName
+	instance.Hostname = conf.Hostname
 	instance.Alias = conf.Alias
 	instance.Cert = conf.Cert
 	instance.Key = conf.Key
diff --git a/pwd/instance_test.go b/pwd/instance_test.go
index 15a567a..b2e0320 100644
--- a/pwd/instance_test.go
+++ b/pwd/instance_test.go
@@ -134,6 +134,54 @@ func TestInstanceNew_WithNotAllowedImage(t *testing.T) {
 	assert.Equal(t, expectedContainerOpts, containerOpts)
 }
 
+func TestInstanceNew_WithCustomHostname(t *testing.T) {
+	containerOpts := docker.CreateContainerOpts{}
+	dock := &mockDocker{}
+	dock.createContainer = func(opts docker.CreateContainerOpts) (string, error) {
+		containerOpts = opts
+		return "10.0.0.1", nil
+	}
+
+	tasks := &mockTasks{}
+	broadcast := &mockBroadcast{}
+	storage := &mockStorage{}
+
+	p := NewPWD(dock, tasks, broadcast, storage)
+
+	session, err := p.SessionNew(time.Hour, "", "")
+
+	assert.Nil(t, err)
+
+	instance, err := p.InstanceNew(session, InstanceConfig{ImageName: "redis", Hostname: "redis-master"})
+
+	assert.Nil(t, err)
+
+	expectedInstance := Instance{
+		Name:         fmt.Sprintf("%s_redis-master", session.Id[:8]),
+		Hostname:     "redis-master",
+		IP:           "10.0.0.1",
+		Alias:        "",
+		Image:        "redis",
+		IsDockerHost: false,
+		session:      session,
+	}
+
+	assert.Equal(t, expectedInstance, *instance)
+
+	expectedContainerOpts := docker.CreateContainerOpts{
+		Image:         expectedInstance.Image,
+		SessionId:     session.Id,
+		PwdIpAddress:  session.PwdIpAddress,
+		ContainerName: expectedInstance.Name,
+		Hostname:      expectedInstance.Hostname,
+		ServerCert:    nil,
+		ServerKey:     nil,
+		CACert:        nil,
+		Privileged:    false,
+	}
+	assert.Equal(t, expectedContainerOpts, containerOpts)
+}
+
 func TestInstanceAllowedImages(t *testing.T) {
 	dock := &mockDocker{}
 	tasks := &mockTasks{}