func (a DrainAction) determineParams(drainType DrainType, currentSpec boshas.V1ApplySpec, newSpecs []boshas.V1ApplySpec) (boshdrain.ScriptParams, error) { var newSpec *boshas.V1ApplySpec var params boshdrain.ScriptParams if len(newSpecs) > 0 { newSpec = &newSpecs[0] } switch drainType { case DrainTypeStatus: // Status was used in the past when dynamic drain was implemented in the Director. // Now that we implement it in the agent, we should never get a call for this type. return params, bosherr.Error("Unexpected call with drain type 'status'") case DrainTypeUpdate: if newSpec == nil { return params, bosherr.Error("Drain update requires new spec") } params = boshdrain.NewUpdateParams(currentSpec, *newSpec) case DrainTypeShutdown: err := a.notifier.NotifyShutdown() if err != nil { return params, bosherr.WrapError(err, "Notifying shutdown") } params = boshdrain.NewShutdownParams(currentSpec, newSpec) } return params, nil }
Expect(err).ToNot(HaveOccurred()) Expect(value).To(Equal(0)) Expect(notifier.NotifiedShutdown).To(BeTrue()) }) Context("when job shutdown notification succeeds", func() { It("runs drain script with shutdown params in parallel", func() { fooScript := &fakescript.FakeCancellableScript{} fooScript.TagReturns("foo") barScript := &fakescript.FakeCancellableScript{} barScript.TagReturns("bar") jobScriptProvider.NewDrainScriptStub = func(jobName string, params boshdrain.ScriptParams) boshscript.CancellableScript { Expect(params).To(Equal(boshdrain.NewShutdownParams(currentSpec, nil))) if jobName == "foo" { return fooScript } else if jobName == "bar" { return barScript } else { panic("Non-matching shutdown drain script created") } } parallelScript.RunReturns(nil) value, err := act() Expect(err).ToNot(HaveOccurred()) Expect(value).To(Equal(0))