func (ts *taskSuite) TestAt(c *C) {
b := new(fakeStateBackend)
b.ensureBefore = time.Hour
st := state.New(b)
st.Lock()
defer st.Unlock()
t := st.NewTask("download", "1...")
now := time.Now()
restore := state.MockTime(now)
defer restore()
when := now.Add(10 * time.Second)
t.At(when)
c.Check(t.AtTime().Equal(when), Equals, true)
c.Check(b.ensureBefore, Equals, 10*time.Second)
}
func (ts *taskRunnerSuite) TestRetryAfterDuration(c *C) {
ensureBeforeTick := make(chan bool, 1)
sb := &stateBackend{
ensureBefore: time.Hour,
ensureBeforeSeen: ensureBeforeTick,
}
st := state.New(sb)
r := state.NewTaskRunner(st)
defer r.Stop()
ch := make(chan bool)
ask := 0
r.AddHandler("ask-for-retry", func(t *state.Task, _ *tomb.Tomb) error {
ask++
if ask == 1 {
return &state.Retry{After: time.Minute}
}
ch <- true
return nil
}, nil)
st.Lock()
chg := st.NewChange("install", "...")
t := st.NewTask("ask-for-retry", "...")
chg.AddTask(t)
st.Unlock()
tock := time.Now()
restore := state.MockTime(tock)
defer restore()
r.Ensure() // will run and be rescheduled in a minute
select {
case <-ensureBeforeTick:
case <-time.After(2 * time.Second):
c.Fatal("EnsureBefore wasn't called")
}
st.Lock()
defer st.Unlock()
c.Check(t.Status(), Equals, state.DoingStatus)
c.Check(ask, Equals, 1)
c.Check(sb.ensureBefore, Equals, 1*time.Minute)
schedule := t.AtTime()
c.Check(schedule.IsZero(), Equals, false)
state.MockTime(tock.Add(5 * time.Second))
sb.ensureBefore = time.Hour
st.Unlock()
r.Ensure() // too soon
st.Lock()
c.Check(t.Status(), Equals, state.DoingStatus)
c.Check(ask, Equals, 1)
c.Check(sb.ensureBefore, Equals, 55*time.Second)
c.Check(t.AtTime().Equal(schedule), Equals, true)
state.MockTime(schedule)
sb.ensureBefore = time.Hour
st.Unlock()
r.Ensure() // time to run again
select {
case <-ch:
case <-time.After(2 * time.Second):
c.Fatal("handler wasn't called")
}
// wait for handler to finish
r.Wait()
st.Lock()
c.Check(t.Status(), Equals, state.DoneStatus)
c.Check(ask, Equals, 2)
c.Check(sb.ensureBefore, Equals, time.Hour)
c.Check(t.AtTime().IsZero(), Equals, true)
}