Added a context prepare step.

worker/uniter/operation/runaction.go

@@ -49,6 +49,10 @@ func (ra *runAction) Prepare(state State) (*State, error) {
 		// this should *really* never happen, but let's not panic
 		return nil, errors.Trace(err)
 	}
+	err = rnr.Context().Prepare()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
 	ra.name = actionData.Name
 	ra.runner = rnr
 	return stateChange{

worker/uniter/operation/runaction_test.go

@@ -99,6 +99,49 @@ func (s *RunActionSuite) TestPrepareErrorOther(c *gc.C) {
 	c.Assert(*runnerFactory.MockNewActionRunner.gotActionId, gc.Equals, someActionId)
 }
 
+func (s *RunActionSuite) TestPrepareCtxCalled(c *gc.C) {
+	ctx := &MockContext{actionData: &runner.ActionData{Name: "some-action-name"}}
+	runnerFactory := &MockRunnerFactory{
+		MockNewActionRunner: &MockNewActionRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
+	}
+	factory := operation.NewFactory(operation.FactoryParams{
+		RunnerFactory: runnerFactory,
+	})
+	op, err := factory.NewAction(someActionId)
+	c.Assert(err, jc.ErrorIsNil)
+
+	newState, err := op.Prepare(operation.State{})
+	c.Assert(err, jc.ErrorIsNil)
+	c.Assert(newState, gc.NotNil)
+	ctx.CheckCall(c, 0, "Prepare")
+}
+
+func (s *RunActionSuite) TestPrepareCtxError(c *gc.C) {
+	ctx := &MockContext{actionData: &runner.ActionData{Name: "some-action-name"}}
+	ctx.SetErrors(errors.New("ctx prepare error"))
+	runnerFactory := &MockRunnerFactory{
+		MockNewActionRunner: &MockNewActionRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
+	}
+	factory := operation.NewFactory(operation.FactoryParams{
+		RunnerFactory: runnerFactory,
+	})
+	op, err := factory.NewAction(someActionId)
+	c.Assert(err, jc.ErrorIsNil)
+
+	newState, err := op.Prepare(operation.State{})
+	c.Assert(err, gc.ErrorMatches, `ctx prepare error`)
+	c.Assert(newState, gc.IsNil)
+	ctx.CheckCall(c, 0, "Prepare")
+}
+
 func (s *RunActionSuite) TestPrepareSuccessCleanState(c *gc.C) {
 	runnerFactory := NewRunActionRunnerFactory(errors.New("should not call"))
 	factory := operation.NewFactory(operation.FactoryParams{

worker/uniter/operation/runcommands.go

@@ -47,7 +47,12 @@ func (rc *runCommands) Prepare(state State) (*State, error) {
 	if err != nil {
 		return nil, err
 	}
+	err = rnr.Context().Prepare()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
 	rc.runner = rnr
+
 	return nil, nil
 }
 

worker/uniter/operation/runcommands_test.go

@@ -42,8 +42,13 @@ func (s *RunCommandsSuite) TestPrepareError(c *gc.C) {
 }
 
 func (s *RunCommandsSuite) TestPrepareSuccess(c *gc.C) {
+	ctx := &MockContext{}
 	runnerFactory := &MockRunnerFactory{
-		MockNewCommandRunner: &MockNewCommandRunner{},
+		MockNewCommandRunner: &MockNewCommandRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
 	}
 	factory := operation.NewFactory(operation.FactoryParams{
 		RunnerFactory: runnerFactory,
@@ -60,6 +65,30 @@ func (s *RunCommandsSuite) TestPrepareSuccess(c *gc.C) {
 		RemoteUnitName:  "foo/456",
 		ForceRemoteUnit: true,
 	})
+	ctx.CheckCall(c, 0, "Prepare")
+}
+
+func (s *RunCommandsSuite) TestPrepareCtxError(c *gc.C) {
+	ctx := &MockContext{}
+	ctx.SetErrors(errors.New("ctx prepare error"))
+	runnerFactory := &MockRunnerFactory{
+		MockNewCommandRunner: &MockNewCommandRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
+	}
+	factory := operation.NewFactory(operation.FactoryParams{
+		RunnerFactory: runnerFactory,
+	})
+	sendResponse := func(*utilexec.ExecResponse, error) { panic("not expected") }
+	op, err := factory.NewCommands(someCommandArgs, sendResponse)
+	c.Assert(err, jc.ErrorIsNil)
+
+	newState, err := op.Prepare(operation.State{})
+	c.Assert(err, gc.ErrorMatches, "ctx prepare error")
+	c.Assert(newState, gc.IsNil)
+	ctx.CheckCall(c, 0, "Prepare")
 }
 
 func (s *RunCommandsSuite) TestExecuteRebootErrors(c *gc.C) {

worker/uniter/operation/runhook.go

@@ -55,6 +55,10 @@ func (rh *runHook) Prepare(state State) (*State, error) {
 	if err != nil {
 		return nil, err
 	}
+	err = rnr.Context().Prepare()
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
 	rh.name = name
 	rh.runner = rnr
 

worker/uniter/operation/runhook_test.go

@@ -77,6 +77,63 @@ func (s *RunHookSuite) testPrepareHookError(
 	})
 }
 
+func (s *RunHookSuite) TestPrepareHookCtxCalled(c *gc.C) {
+	ctx := &MockContext{}
+	callbacks := &PrepareHookCallbacks{
+		MockPrepareHook:       &MockPrepareHook{},
+		MockClearResolvedFlag: &MockNoArgs{},
+	}
+	runnerFactory := &MockRunnerFactory{
+		MockNewHookRunner: &MockNewHookRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
+	}
+	factory := operation.NewFactory(operation.FactoryParams{
+		RunnerFactory: runnerFactory,
+		Callbacks:     callbacks,
+	})
+
+	op, err := factory.NewRunHook(hook.Info{Kind: hooks.ConfigChanged})
+	c.Assert(err, jc.ErrorIsNil)
+
+	newState, err := op.Prepare(operation.State{})
+	c.Check(newState, gc.NotNil)
+	c.Assert(err, jc.ErrorIsNil)
+
+	ctx.CheckCall(c, 0, "Prepare")
+}
+
+func (s *RunHookSuite) TestPrepareHookCtxError(c *gc.C) {
+	ctx := &MockContext{}
+	ctx.SetErrors(errors.New("ctx prepare error"))
+	callbacks := &PrepareHookCallbacks{
+		MockPrepareHook:       &MockPrepareHook{},
+		MockClearResolvedFlag: &MockNoArgs{},
+	}
+	runnerFactory := &MockRunnerFactory{
+		MockNewHookRunner: &MockNewHookRunner{
+			runner: &MockRunner{
+				context: ctx,
+			},
+		},
+	}
+	factory := operation.NewFactory(operation.FactoryParams{
+		RunnerFactory: runnerFactory,
+		Callbacks:     callbacks,
+	})
+
+	op, err := factory.NewRunHook(hook.Info{Kind: hooks.ConfigChanged})
+	c.Assert(err, jc.ErrorIsNil)
+
+	newState, err := op.Prepare(operation.State{})
+	c.Check(newState, gc.IsNil)
+	c.Assert(err, gc.ErrorMatches, `ctx prepare error`)
+
+	ctx.CheckCall(c, 0, "Prepare")
+}
+
 func (s *RunHookSuite) TestPrepareHookError_Run(c *gc.C) {
 	s.testPrepareHookError(c, (operation.Factory).NewRunHook, false, false)
 }

worker/uniter/operation/util_test.go

@@ -5,6 +5,7 @@ package operation_test
 
 import (
 	"github.com/juju/errors"
+	"github.com/juju/testing"
 	utilexec "github.com/juju/utils/exec"
 	corecharm "gopkg.in/juju/charm.v5"
 	"gopkg.in/juju/charm.v5/hooks"
@@ -294,6 +295,7 @@ func (f *MockRunnerFactory) NewCommandRunner(commandInfo runner.CommandInfo) (ru
 
 type MockContext struct {
 	runner.Context
+	testing.Stub
 	actionData      *runner.ActionData
 	setStatusCalled bool
 	status          jujuc.StatusInfo
@@ -324,6 +326,11 @@ func (mock *MockContext) UnitStatus() (*jujuc.StatusInfo, error) {
 	return &mock.status, nil
 }
 
+func (mock *MockContext) Prepare() error {
+	mock.MethodCall(mock, "Prepare")
+	return mock.NextErr()
+}
+
 type MockRunAction struct {
 	gotName *string
 	err     error
@@ -427,6 +434,7 @@ func NewRunCommandsRunnerFactory(runResponse *utilexec.ExecResponse, runErr erro
 		MockNewCommandRunner: &MockNewCommandRunner{
 			runner: &MockRunner{
 				MockRunCommands: &MockRunCommands{response: runResponse, err: runErr},
+				context:         &MockContext{},
 			},
 		},
 	}

worker/uniter/runner/context.go

@@ -570,6 +570,17 @@ func (ctx *HookContext) addJujuUnitsMetric() error {
 	return nil
 }
 
+// Prepare implements the Context interface.
+func (ctx *HookContext) Prepare() error {
+	if ctx.actionData != nil {
+		err := ctx.state.ActionBegin(ctx.actionData.Tag)
+		if err != nil {
+			return errors.Trace(err)
+		}
+	}
+	return nil
+}
+
 // Flush implements the Context interface.
 func (ctx *HookContext) Flush(process string, ctxErr error) (err error) {
 	// A non-existant metricsRecorder simply means that metrics were disabled

worker/uniter/runner/factory.go

@@ -242,11 +242,6 @@ func (f *factory) NewActionRunner(actionId string) (Runner, error) {
 		return nil, errors.Trace(err)
 	}
 
-	err = f.state.ActionBegin(tag)
-	if err != nil {
-		return nil, errors.Trace(err)
-	}
-
 	name := action.Name()
 	spec, ok := ch.Actions().ActionSpecs[name]
 	if !ok {

worker/uniter/runner/runner.go

@@ -42,9 +42,11 @@ type Context interface {
 	HookVars(paths Paths) []string
 	ActionData() (*ActionData, error)
 	SetProcess(process *os.Process)
-	Flush(badge string, failure error) error
 	HasExecutionSetUnitStatus() bool
 	ResetExecutionSetUnitStatus()
+
+	Prepare() error
+	Flush(badge string, failure error) error
 }
 
 // Paths exposes the paths needed by Runner.
@@ -97,6 +99,7 @@ func (runner *runner) RunCommands(commands string) (*utilexec.ExecResponse, erro
 		WorkingDir:  runner.paths.GetCharmDir(),
 		Environment: env,
 	}
+
 	err = command.Run()
 	if err != nil {
 		return nil, err

worker/uniter/runner/runner_test.go

@@ -172,6 +172,10 @@ func (ctx *MockContext) SetProcess(process *os.Process) {
 	ctx.expectPid = process.Pid
 }
 
+func (ctx *MockContext) Prepare() error {
+	return nil
+}
+
 func (ctx *MockContext) Flush(badge string, failure error) error {
 	ctx.flushBadge = badge
 	ctx.flushFailure = failure