From c9c294f1361348270da60099b280e65c4f8ca0aa Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 7 May 2013 17:10:45 -0700
Subject: [PATCH] packer/plugin: More robust command exit detection + tests

---
 packer/plugin/command.go      | 25 +++++++++++++++++--------
 packer/plugin/command_test.go |  9 +++++++++
 packer/plugin/plugin_test.go  |  4 ++++
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/packer/plugin/command.go b/packer/plugin/command.go
index 3f5a66b0e..beb677e7d 100644
--- a/packer/plugin/command.go
+++ b/packer/plugin/command.go
@@ -36,21 +36,30 @@ func Command(cmd *exec.Cmd) (result packer.Command, err error) {
 		select {
 		case <-cmdExited:
 			err = errors.New("plugin exited before we could connect")
-			return
-		case <-time.After(10 * time.Millisecond):
-			if line, err := out.ReadBytes('\n'); err == nil {
-				address = strings.TrimSpace(string(line))
-				done = true
-			}
+			done = true
+		default:
+		}
 
-			// Make sure to reset err to nil
+		if line, lerr := out.ReadBytes('\n'); lerr == nil {
+			// Trim the address and reset the err since we were able
+			// to read some sort of address.
+			address = strings.TrimSpace(string(line))
 			err = nil
+			break
 		}
+
+		// If error is nil from previously, return now
+		if err != nil {
+			return
+		}
+
+		// Wait a bit
+		time.Sleep(10 * time.Millisecond)
 	}
 
 	client, err := rpc.Dial("tcp", address)
 	if err != nil {
-		panic(err)
+		return
 	}
 
 	result = packrpc.Command(client)
diff --git a/packer/plugin/command_test.go b/packer/plugin/command_test.go
index 9dfe56464..8404661f5 100644
--- a/packer/plugin/command_test.go
+++ b/packer/plugin/command_test.go
@@ -28,4 +28,13 @@ func TestCommand_CommandExited(t *testing.T) {
 
 	_, err := Command(helperProcess("im-a-command-that-doesnt-work"))
 	assert.NotNil(err, "should have an error")
+	assert.Equal(err.Error(), "plugin exited before we could connect", "be correct error")
+}
+
+func TestCommand_BadRPC(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	_, err := Command(helperProcess("invalid-rpc-address"))
+	assert.NotNil(err, "should have an error")
+	assert.Equal(err.Error(), "missing port in address lolinvalid", "be correct error")
 }
diff --git a/packer/plugin/plugin_test.go b/packer/plugin/plugin_test.go
index e8397085e..35cb84145 100644
--- a/packer/plugin/plugin_test.go
+++ b/packer/plugin/plugin_test.go
@@ -40,6 +40,8 @@ func TestHelperProcess(*testing.T) {
 		return
 	}
 
+	defer os.Exit(0)
+
 	args := os.Args
 	for len(args) > 0 {
 		if args[0] == "--" {
@@ -59,6 +61,8 @@ func TestHelperProcess(*testing.T) {
 	switch cmd {
 	case "command":
 		ServeCommand(new(helperCommand))
+	case "invalid-rpc-address":
+		fmt.Println("lolinvalid")
 	case "start-timeout":
 		time.Sleep(1 * time.Minute)
 		os.Exit(1)