diff --git a/packer/hook_mock.go b/packer/hook_mock.go index eedf5bb90..7177329e3 100644 --- a/packer/hook_mock.go +++ b/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() { diff --git a/packer/rpc/hook_test.go b/packer/rpc/hook_test.go index 9dc53aeac..1d226056a 100644 --- a/packer/rpc/hook_test.go +++ b/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) + } +}