packer/rpc: test concurrent cancel/run

packer/hook_mock.go

@@ -2,6 +2,8 @@ package packer
 
 // MockHook is an implementation of Hook that can be used for tests.
 type MockHook struct {
+	RunFunc func() error
+
 	RunCalled    bool
 	RunComm      Communicator
 	RunData      interface{}
@@ -16,7 +18,12 @@ func (t *MockHook) Run(name string, ui Ui, comm Communicator, data interface{})
 	t.RunData = data
 	t.RunName = name
 	t.RunUi = ui
-	return nil
+
+	if t.RunFunc == nil {
+		return nil
+	}
+
+	return t.RunFunc()
 }
 
 func (t *MockHook) Cancel() {

packer/rpc/hook_test.go

@@ -4,7 +4,10 @@ import (
 	"cgl.tideland.biz/asserts"
 	"github.com/mitchellh/packer/packer"
 	"net/rpc"
+	"reflect"
+	"sync"
 	"testing"
+	"time"
 )
 
 func TestHookRPC(t *testing.T) {
@@ -42,3 +45,56 @@ func TestHook_Implements(t *testing.T) {
 
 	assert.Implementor(h, &r, "should be a Hook")
 }
+
+func TestHook_cancelWhileRun(t *testing.T) {
+	var finishLock sync.Mutex
+	finishOrder := make([]string, 0, 2)
+
+	h := &packer.MockHook{
+		RunFunc: func() error {
+			time.Sleep(100 * time.Millisecond)
+
+			finishLock.Lock()
+			finishOrder = append(finishOrder, "run")
+			finishLock.Unlock()
+			return nil
+		},
+	}
+
+	// Serve
+	server := rpc.NewServer()
+	RegisterHook(server, h)
+	address := serveSingleConn(server)
+
+	// Create the client over RPC and run some methods to verify it works
+	client, err := rpc.Dial("tcp", address)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	hClient := Hook(client)
+
+	// Start the run
+	finished := make(chan struct{})
+	go func() {
+		hClient.Run("foo", nil, nil, nil)
+		close(finished)
+	}()
+
+	// Cancel it pretty quickly.
+	time.Sleep(10 * time.Millisecond)
+	hClient.Cancel()
+
+	finishLock.Lock()
+	finishOrder = append(finishOrder, "cancel")
+	finishLock.Unlock()
+
+	// Verify things are good
+	<-finished
+
+	// Check the results
+	expected := []string{"cancel", "run"}
+	if !reflect.DeepEqual(finishOrder, expected) {
+		t.Fatalf("bad: %#v", finishOrder)
+	}
+}