FIX docker#2457 Add a flag to set the CPU count

Signed-off-by: David Gageot <david@gageot.net>
dgageot · Dec 23, 2015 · afcaaae · afcaaae
1 parent 63ca58d
commit afcaaae
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 4 deletions.
diff --git a/drivers/hyperv/hyperv.go b/drivers/hyperv/hyperv.go
@@ -24,17 +24,20 @@ type Driver struct {
 	diskImage      string
 	DiskSize       int
 	MemSize        int
+	CPU            int
 }
 
 const (
 	defaultDiskSize = 20000
 	defaultMemory   = 1024
+	defaultCPU      = 1
 )
 
-func NewDriver(hostName, storePath string) drivers.Driver {
+func NewDriver(hostName, storePath string) *Driver {
 	return &Driver{
 		DiskSize: defaultDiskSize,
 		MemSize:  defaultMemory,
+		CPU:      defaultCPU,
 		BaseDriver: &drivers.BaseDriver{
 			MachineName: hostName,
 			StorePath:   storePath,
@@ -64,6 +67,11 @@ func (d *Driver) GetCreateFlags() []mcnflag.Flag {
 			Usage: "Memory size for host in MB.",
 			Value: defaultMemory,
 		},
+		mcnflag.IntFlag{
+			Name:  "hyperv-cpu-count",
+			Usage: "number of CPUs for the machine",
+			Value: defaultCPU,
+		},
 	}
 }
 
@@ -72,6 +80,7 @@ func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 	d.vSwitch = flags.String("hyperv-virtual-switch")
 	d.DiskSize = flags.Int("hyperv-disk-size")
 	d.MemSize = flags.Int("hyperv-memory")
+	d.CPU = flags.Int("hyperv-cpu-count")
 	d.SwarmMaster = flags.Bool("swarm-master")
 	d.SwarmHost = flags.String("swarm-host")
 	d.SwarmDiscovery = flags.String("swarm-discovery")
@@ -181,6 +190,12 @@ func (d *Driver) Create() error {
 		return err
 	}
 
+	if d.CPU > 1 {
+		if err := cmd("SET-VMProcessor", "-Name", d.MachineName, "-Count", fmt.Sprintf("%d", d.CPU)); err != nil {
+			return err
+		}
+	}
+
 	if err := cmd("Set-VMDvdDrive", "-VMName", d.MachineName, "-Path", fmt.Sprintf("'%s'", d.ResolveStorePath("boot2docker.iso"))); err != nil {
 		return err
 	}

diff --git a/drivers/hyperv/hyperv_test.go b/drivers/hyperv/hyperv_test.go
@@ -7,7 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestSetConfigFromFlags(t *testing.T) {
+func TestSetConfigFromDefaultFlags(t *testing.T) {
 	driver := NewDriver("default", "path")
 
 	checkFlags := &drivers.CheckDriverOptions{
@@ -19,10 +19,44 @@ func TestSetConfigFromFlags(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Empty(t, checkFlags.InvalidFlags)
 
-	sshUser := driver.GetSSHUsername()
-	assert.Equal(t, "docker", sshUser)
+	sshPort, err := driver.GetSSHPort()
+	assert.Equal(t, 22, sshPort)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "", driver.Boot2DockerURL)
+	assert.Equal(t, "", driver.vSwitch)
+	assert.Equal(t, defaultDiskSize, driver.DiskSize)
+	assert.Equal(t, defaultMemory, driver.MemSize)
+	assert.Equal(t, defaultCPU, driver.CPU)
+	assert.Equal(t, "docker", driver.GetSSHUsername())
+}
+
+func TestSetConfigFromCustomFlags(t *testing.T) {
+	driver := NewDriver("default", "path")
+
+	checkFlags := &drivers.CheckDriverOptions{
+		FlagsValues: map[string]interface{}{
+			"hyperv-boot2docker-url": "B2D_URL",
+			"hyperv-virtual-switch":  "TheSwitch",
+			"hyperv-disk-size":       100000,
+			"hyperv-memory":          4096,
+			"hyperv-cpu-count":       4,
+		},
+		CreateFlags: driver.GetCreateFlags(),
+	}
+
+	err := driver.SetConfigFromFlags(checkFlags)
+	assert.NoError(t, err)
+	assert.Empty(t, checkFlags.InvalidFlags)
 
 	sshPort, err := driver.GetSSHPort()
 	assert.Equal(t, 22, sshPort)
 	assert.NoError(t, err)
+
+	assert.Equal(t, "B2D_URL", driver.Boot2DockerURL)
+	assert.Equal(t, "TheSwitch", driver.vSwitch)
+	assert.Equal(t, 100000, driver.DiskSize)
+	assert.Equal(t, 4096, driver.MemSize)
+	assert.Equal(t, 4, driver.CPU)
+	assert.Equal(t, "docker", driver.GetSSHUsername())
 }