From f890f9b6cd3e28071b353bee84144e7c0b7bb6d9 Mon Sep 17 00:00:00 2001
From: Nathan LeClaire <nathan.leclaire@gmail.com>
Date: Tue, 10 Mar 2015 16:36:13 -0700
Subject: [PATCH] Implement wait for desired state before exiting

Signed-off-by: Nathan LeClaire <nathan.leclaire@gmail.com>
---
 commands.go |  8 ++++----
 host.go     | 47 +++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/commands.go b/commands.go
index bbdead7a16..9d1da8f70a 100644
--- a/commands.go
+++ b/commands.go
@@ -574,10 +574,10 @@ func cmdSsh(c *cli.Context) {
 // We run commands concurrently and communicate back an error if there was one.
 func machineCommand(actionName string, machine *Host, errorChan chan<- error) {
 	commands := map[string](func() error){
-		"start":   machine.Driver.Start,
-		"stop":    machine.Driver.Stop,
-		"restart": machine.Driver.Restart,
-		"kill":    machine.Driver.Kill,
+		"start":   machine.Start,
+		"stop":    machine.Stop,
+		"restart": machine.Restart,
+		"kill":    machine.Kill,
 		"upgrade": machine.Upgrade,
 	}
 
diff --git a/host.go b/host.go
index d579599e7f..2546101527 100644
--- a/host.go
+++ b/host.go
@@ -21,6 +21,7 @@ import (
 	"github.com/docker/machine/drivers"
 	"github.com/docker/machine/provider"
 	"github.com/docker/machine/ssh"
+	"github.com/docker/machine/state"
 	"github.com/docker/machine/utils"
 )
 
@@ -588,12 +589,54 @@ func (h *Host) SetHostname() error {
 	return nil
 }
 
+func (h *Host) MachineInState(desiredState state.State) func() bool {
+	return func() bool {
+		currentState, err := h.Driver.GetState()
+		if err != nil {
+			log.Debugf("Error getting machine state: %s", err)
+		}
+		if currentState == desiredState {
+			return true
+		}
+		return false
+	}
+}
+
 func (h *Host) Start() error {
-	return h.Driver.Start()
+	if err := h.Driver.Start(); err != nil {
+		return err
+	}
+	return utils.WaitFor(h.MachineInState(state.Running))
 }
 
 func (h *Host) Stop() error {
-	return h.Driver.Stop()
+	if err := h.Driver.Stop(); err != nil {
+		return err
+	}
+	return utils.WaitFor(h.MachineInState(state.Stopped))
+}
+
+func (h *Host) Kill() error {
+	if err := h.Driver.Stop(); err != nil {
+		return err
+	}
+	return utils.WaitFor(h.MachineInState(state.Stopped))
+}
+
+func (h *Host) Restart() error {
+	if err := h.Stop(); err != nil {
+		return err
+	}
+	if err := utils.WaitFor(h.MachineInState(state.Stopped)); err != nil {
+		return err
+	}
+	if err := h.Start(); err != nil {
+		return err
+	}
+	if err := utils.WaitFor(h.MachineInState(state.Running)); err != nil {
+		return err
+	}
+	return nil
 }
 
 func (h *Host) Upgrade() error {