builder/vmware/iso: Move remote registration out to separate step

builder/vmware/iso/builder.go

@@ -395,6 +395,7 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		&vmwcommon.StepSuppressMessages{},
 		&stepHTTPServer{},
 		&stepConfigureVNC{},
+		&StepRegister{},
 		&stepRun{},
 		&stepTypeBootCommand{},
 		&common.StepConnectSSH{

builder/vmware/iso/remote_driver_mock.go

+package iso
+
+import (
+	vmwcommon "github.com/mitchellh/packer/builder/vmware/common"
+)
+
+type RemoteDriverMock struct {
+	vmwcommon.DriverMock
+
+	UploadISOCalled bool
+	UploadISOPath   string
+	UploadISOResult string
+	UploadISOErr    error
+
+	RegisterCalled bool
+	RegisterPath   string
+	RegisterErr    error
+
+	UnregisterCalled bool
+	UnregisterPath   string
+	UnregisterErr    error
+}
+
+func (d *RemoteDriverMock) UploadISO(path string) (string, error) {
+	d.UploadISOCalled = true
+	d.UploadISOPath = path
+	return d.UploadISOResult, d.UploadISOErr
+}
+
+func (d *RemoteDriverMock) Register(path string) error {
+	d.RegisterCalled = true
+	d.RegisterPath = path
+	return d.RegisterErr
+}
+
+func (d *RemoteDriverMock) Unregister(path string) error {
+	d.UnregisterCalled = true
+	d.UnregisterPath = path
+	return d.UnregisterErr
+}

builder/vmware/iso/remote_driver_mock_test.go

+package iso
+
+import (
+	"testing"
+
+	vmwcommon "github.com/mitchellh/packer/builder/vmware/common"
+)
+
+func TestRemoteDriverMock_impl(t *testing.T) {
+	var _ vmwcommon.Driver = new(RemoteDriverMock)
+	var _ RemoteDriver = new(RemoteDriverMock)
+}

builder/vmware/iso/step_register.go

+package iso
+
+import (
+	"fmt"
+
+	"github.com/mitchellh/multistep"
+	vmwcommon "github.com/mitchellh/packer/builder/vmware/common"
+	"github.com/mitchellh/packer/packer"
+)
+
+type StepRegister struct {
+	registeredPath string
+}
+
+func (s *StepRegister) Run(state multistep.StateBag) multistep.StepAction {
+	driver := state.Get("driver").(vmwcommon.Driver)
+	ui := state.Get("ui").(packer.Ui)
+	vmxPath := state.Get("vmx_path").(string)
+
+	if remoteDriver, ok := driver.(RemoteDriver); ok {
+		ui.Say("Registering remote VM...")
+		if err := remoteDriver.Register(vmxPath); err != nil {
+			err := fmt.Errorf("Error registering VM: %s", err)
+			state.Put("error", err)
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+
+		s.registeredPath = vmxPath
+	}
+
+	return multistep.ActionContinue
+}
+
+func (s *StepRegister) Cleanup(state multistep.StateBag) {
+	if s.registeredPath == "" {
+		return
+	}
+
+	driver := state.Get("driver").(vmwcommon.Driver)
+	ui := state.Get("ui").(packer.Ui)
+
+	if remoteDriver, ok := driver.(RemoteDriver); ok {
+		ui.Say("Unregistering virtual machine...")
+		if err := remoteDriver.Unregister(s.registeredPath); err != nil {
+			ui.Error(fmt.Sprintf("Error unregistering VM: %s", err))
+		}
+
+		s.registeredPath = ""
+	}
+
+}

builder/vmware/iso/step_register_test.go

+package iso
+
+import (
+	"github.com/mitchellh/multistep"
+	"testing"
+)
+
+func TestStepRegister_impl(t *testing.T) {
+	var _ multistep.Step = new(StepRegister)
+}
+
+func TestStepRegister_regularDriver(t *testing.T) {
+	state := testState(t)
+	step := new(StepRegister)
+
+	state.Put("vmx_path", "foo")
+
+	// Test the run
+	if action := step.Run(state); action != multistep.ActionContinue {
+		t.Fatalf("bad action: %#v", action)
+	}
+	if _, ok := state.GetOk("error"); ok {
+		t.Fatal("should NOT have error")
+	}
+
+	// Cleanup
+	step.Cleanup(state)
+}
+
+func TestStepRegister_remoteDriver(t *testing.T) {
+	state := testState(t)
+	step := new(StepRegister)
+
+	driver := new(RemoteDriverMock)
+	state.Put("driver", driver)
+	state.Put("vmx_path", "foo")
+
+	// Test the run
+	if action := step.Run(state); action != multistep.ActionContinue {
+		t.Fatalf("bad action: %#v", action)
+	}
+	if _, ok := state.GetOk("error"); ok {
+		t.Fatal("should NOT have error")
+	}
+
+	// verify
+	if !driver.RegisterCalled {
+		t.Fatal("register should be called")
+	}
+	if driver.RegisterPath != "foo" {
+		t.Fatal("should call with correct path")
+	}
+	if driver.UnregisterCalled {
+		t.Fatal("unregister should not be called")
+	}
+
+	// cleanup
+	step.Cleanup(state)
+	if !driver.UnregisterCalled {
+		t.Fatal("unregister should be called")
+	}
+	if driver.UnregisterPath != "foo" {
+		t.Fatal("should unregister proper path")
+	}
+}

builder/vmware/iso/step_run.go

@@ -21,8 +21,6 @@ import (
 type stepRun struct {
 	bootTime time.Time
 	vmxPath  string
-
-	registered bool
 }
 
 func (s *stepRun) Run(state multistep.StateBag) multistep.StepAction {
@@ -45,17 +43,6 @@ func (s *stepRun) Run(state multistep.StateBag) multistep.StepAction {
 				"%s:%d", vncIp, vncPort))
 	}
 
-	if remoteDriver, ok := driver.(RemoteDriver); ok {
-		if err := remoteDriver.Register(vmxPath); err != nil {
-			err := fmt.Errorf("Error registering VM: %s", err)
-			state.Put("error", err)
-			ui.Error(err.Error())
-			return multistep.ActionHalt
-		}
-
-		s.registered = true
-	}
-
 	if err := driver.Start(vmxPath, config.Headless); err != nil {
 		err := fmt.Errorf("Error starting VM: %s", err)
 		state.Put("error", err)
@@ -107,14 +94,5 @@ func (s *stepRun) Cleanup(state multistep.StateBag) {
 				ui.Error(fmt.Sprintf("Error stopping VM: %s", err))
 			}
 		}
-
-		if remoteDriver, ok := driver.(RemoteDriver); ok && s.registered {
-			ui.Say("Unregistering virtual machine...")
-			if err := remoteDriver.Unregister(s.vmxPath); err != nil {
-				ui.Error(fmt.Sprintf("Error unregistering VM: %s", err))
-			}
-
-			s.registered = false
-		}
 	}
 }

builder/vmware/iso/step_test.go

@@ -3,12 +3,14 @@ package iso
 import (
 	"bytes"
 	"github.com/mitchellh/multistep"
+	vmwcommon "github.com/mitchellh/packer/builder/vmware/common"
 	"github.com/mitchellh/packer/packer"
 	"testing"
 )
 
 func testState(t *testing.T) multistep.StateBag {
 	state := new(multistep.BasicStateBag)
+	state.Put("driver", new(vmwcommon.DriverMock))
 	state.Put("ui", &packer.BasicUi{
 		Reader: new(bytes.Buffer),
 		Writer: new(bytes.Buffer),