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) }