|
5 | 5 | "errors" |
6 | 6 | "slices" |
7 | 7 | "sync" |
| 8 | +"testing" |
8 | 9 | "time" |
9 | 10 | ) |
10 | 11 |
|
@@ -141,14 +142,50 @@ func (m *Mock) matchCallLocked(c *Call) { |
141 | 142 | m.mu.Lock() |
142 | 143 | } |
143 | 144 |
|
144 | | -// Advance moves the clock forward by d, triggering any timers or tickers. Advance will wait for |
145 | | -// tick functions of tickers created using TickerFunc to complete before returning from |
146 | | -// Advance. If multiple timers or tickers trigger simultaneously, they are all run on separate go |
147 | | -// routines. |
148 | | -func (m*Mock)Advance(d time.Duration) { |
149 | | -m.mu.Lock() |
150 | | -deferm.mu.Unlock() |
151 | | -m.advanceLocked(d) |
| 145 | +// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for tick functions of |
| 146 | +// tickers created using TickerFunc to complete. If multiple timers or tickers trigger |
| 147 | +// simultaneously, they are all run on separate go routines. |
| 148 | +typeAdvanceWaiterstruct { |
| 149 | +chchanstruct{} |
| 150 | +} |
| 151 | + |
| 152 | +// Wait for all timers and ticks to complete, or until context expires. |
| 153 | +func (wAdvanceWaiter)Wait(ctx context.Context)error { |
| 154 | +select { |
| 155 | +case<-w.ch: |
| 156 | +returnnil |
| 157 | +case<-ctx.Done(): |
| 158 | +returnctx.Err() |
| 159 | +} |
| 160 | +} |
| 161 | + |
| 162 | +// MustWait waits for all timers and ticks to complete, and fails the test immediately if the |
| 163 | +// context completes first. |
| 164 | +func (wAdvanceWaiter)MustWait(ctx context.Context,t testing.TB) { |
| 165 | +select { |
| 166 | +case<-w.ch: |
| 167 | +return |
| 168 | +case<-ctx.Done(): |
| 169 | +t.Fatalf("context expired waiting for Advance: %s",ctx.Err()) |
| 170 | +} |
| 171 | +} |
| 172 | + |
| 173 | +// Done returns a channel that is closed when all timers and ticks complete. |
| 174 | +func (wAdvanceWaiter)Done()<-chanstruct{} { |
| 175 | +returnw.ch |
| 176 | +} |
| 177 | + |
| 178 | +// Advance moves the clock forward by d, triggering any timers or tickers. The returned value can |
| 179 | +// be used to wait for all timers and ticks to complete. |
| 180 | +func (m*Mock)Advance(d time.Duration)AdvanceWaiter { |
| 181 | +w:=AdvanceWaiter{ch:make(chanstruct{})} |
| 182 | +gofunc() { |
| 183 | +deferclose(w.ch) |
| 184 | +m.mu.Lock() |
| 185 | +deferm.mu.Unlock() |
| 186 | +m.advanceLocked(d) |
| 187 | +}() |
| 188 | +returnw |
152 | 189 | } |
153 | 190 |
|
154 | 191 | func (m*Mock)advanceLocked(d time.Duration) { |
@@ -194,19 +231,24 @@ func (m *Mock) advanceLocked(d time.Duration) { |
194 | 231 | // Set the time to t. If the time is after the current mocked time, then this is equivalent to |
195 | 232 | // Advance() with the difference. You may only Set the time earlier than the current time before |
196 | 233 | // starting tickers and timers (e.g. at the start of your test case). |
197 | | -func (m*Mock)Set(t time.Time) { |
198 | | -m.mu.Lock() |
199 | | -deferm.mu.Unlock() |
200 | | -ift.Before(m.cur) { |
201 | | -// past |
202 | | -if!m.nextTime.IsZero() { |
203 | | -panic("Set mock clock to the past after timers/tickers started") |
| 234 | +func (m*Mock)Set(t time.Time)AdvanceWaiter { |
| 235 | +w:=AdvanceWaiter{ch:make(chanstruct{})} |
| 236 | +gofunc() { |
| 237 | +deferclose(w.ch) |
| 238 | +m.mu.Lock() |
| 239 | +deferm.mu.Unlock() |
| 240 | +ift.Before(m.cur) { |
| 241 | +// past |
| 242 | +if!m.nextTime.IsZero() { |
| 243 | +panic("Set mock clock to the past after timers/tickers started") |
| 244 | +} |
| 245 | +m.cur=t |
| 246 | +return |
204 | 247 | } |
205 | | -m.cur=t |
206 | | -return |
207 | | -} |
208 | | -// future, just advance as normal. |
209 | | -m.advanceLocked(t.Sub(m.cur)) |
| 248 | +// future, just advance as normal. |
| 249 | +m.advanceLocked(t.Sub(m.cur)) |
| 250 | +}() |
| 251 | +returnw |
210 | 252 | } |
211 | 253 |
|
212 | 254 | // Trapper allows the creation of Traps |
|