package quartz

import (

"context"

"errors"

"fmt"

"slices"

"sync"

"time"

)

// TestingT is the minimal interface required from a testing framework for the Mock.

type TestingT interface {

Helper()

Log(...any)

Logf(string, ...any)

Error(...any)

Errorf(string, ...any)

Fatal(...any)

Fatalf(string, ...any)

Cleanup(func())

}

// Mock is the testing implementation of Clock. It tracks a time that monotonically increases

// during a test, triggering any timers or tickers automatically.

type Mock struct {

tb TestingT

logger Logger

mu sync.Mutex

testOver bool

// cur is the current time

cur time.Time

all []event

nextTime time.Time

nextEvents []event

traps []*Trap

}

type event interface {

next() time.Time

fire(t time.Time)

}

func (m *Mock) TickerFunc(ctx context.Context, d time.Duration, f func() error, tags ...string) Waiter {

if d <= 0 {

panic("TickerFunc called with negative or zero duration")

}

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionTickerFunc, tags, withDuration(d))

m.matchCallLocked(c)

defer close(c.complete)

t := &mockTickerFunc{

ctx: ctx,

d: d,

f: f,

nxt: m.cur.Add(d),

mock: m,

cond: sync.NewCond(&m.mu),

}

m.all = append(m.all, t)

m.recomputeNextLocked()

go t.waitForCtx()

return t

}

// NewTicker creates a mocked ticker attached to this Mock. Note that it will cease sending ticks on its channel at the

// end of the test, to avoid leaking any goroutines. Ticks are suppressed even if the mock clock is advanced after the

// test completes. Best practice is to only manipulate the mock time in the main goroutine of the test.

func (m *Mock) NewTicker(d time.Duration, tags ...string) *Ticker {

if d <= 0 {

panic("NewTicker called with negative or zero duration")

}

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionNewTicker, tags, withDuration(d))

m.matchCallLocked(c)

defer close(c.complete)

return newMockTickerLocked(m, d)

}

func (m *Mock) NewTimer(d time.Duration, tags ...string) *Timer {

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionNewTimer, tags, withDuration(d))

defer close(c.complete)

m.matchCallLocked(c)

ch := make(chan time.Time)

t := &Timer{

C: ch,

c: ch,

nxt: m.cur.Add(d),

mock: m,

}

if d <= 0 {

// zero or negative duration timer means we should immediately fire

// it, rather than add it.

go t.fire(t.mock.cur)

return t

}

m.addEventLocked(t)

return t

}

func (m *Mock) AfterFunc(d time.Duration, f func(), tags ...string) *Timer {

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionAfterFunc, tags, withDuration(d))

defer close(c.complete)

m.matchCallLocked(c)

t := &Timer{

nxt: m.cur.Add(d),

fn: f,

mock: m,

}

if d <= 0 {

// zero or negative duration timer means we should immediately fire

// it, rather than add it.

go t.fire(t.mock.cur)

return t

}

m.addEventLocked(t)

return t

}

func (m *Mock) Now(tags ...string) time.Time {

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionNow, tags)

defer close(c.complete)

m.matchCallLocked(c)

return m.cur

}

func (m *Mock) Since(t time.Time, tags ...string) time.Duration {

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionSince, tags, withTime(t))

defer close(c.complete)

m.matchCallLocked(c)

return m.cur.Sub(t)

}

func (m *Mock) Until(t time.Time, tags ...string) time.Duration {

m.mu.Lock()

defer m.mu.Unlock()

c := newCall(clockFunctionUntil, tags, withTime(t))

defer close(c.complete)

m.matchCallLocked(c)

return t.Sub(m.cur)

}

func (m *Mock) addEventLocked(e event) {

m.all = append(m.all, e)

m.recomputeNextLocked()

}

func (m *Mock) recomputeNextLocked() {

var best time.Time

var events []event

for _, e := range m.all {

if best.IsZero() || e.next().Before(best) {

best = e.next()

events = []event{e}

continue

}

if e.next().Equal(best) {

events = append(events, e)

continue

}

}

m.nextTime = best

m.nextEvents = events

}

func (m *Mock) removeTimer(t *Timer) {

m.mu.Lock()

defer m.mu.Unlock()

m.removeTimerLocked(t)

}

func (m *Mock) removeTimerLocked(t *Timer) {

t.stopped = true

m.removeEventLocked(t)

}

func (m *Mock) removeEventLocked(e event) {

defer m.recomputeNextLocked()

for i := range m.all {

if m.all[i] == e {

m.all = append(m.all[:i], m.all[i+1:]...)

return

}

}

}

func (m *Mock) matchCallLocked(c *apiCall) {

var traps []*Trap

for _, t := range m.traps {

if t.matches(c) {

traps = append(traps, t)

}

}

if !m.testOver {

m.logger.Logf("Mock Clock - %s call, matched %d traps", c, len(traps))

}

if len(traps) == 0 {

return

}

c.releases.Add(len(traps))

m.mu.Unlock()

for _, t := range traps {

go t.catch(c)

}

c.releases.Wait()

m.mu.Lock()

}

// AdvanceWaiter is returned from Advance and Set calls and allows you to wait for ticks and timers

// to complete. In the case of functions passed to AfterFunc or TickerFunc, it waits for the

// functions to return. For other ticks & timers, it just waits for the tick to be delivered to

// the channel.

//

// If multiple timers or tickers trigger simultaneously, they are all run on separate

// go routines.

type AdvanceWaiter struct {

tb TestingT

ch chan struct{}

}

// Wait for all timers and ticks to complete, or until context expires.

func (w AdvanceWaiter) Wait(ctx context.Context) error {

select {

case <-w.ch:

return nil

case <-ctx.Done():

return ctx.Err()

}

}

// MustWait waits for all timers and ticks to complete, and fails the test immediately if the

// context completes first. MustWait must be called from the goroutine running the test or

// benchmark, similar to `t.FailNow()`.

func (w AdvanceWaiter) MustWait(ctx context.Context) {

w.tb.Helper()

select {

case <-w.ch:

return

case <-ctx.Done():

w.tb.Fatalf("context expired while waiting for clock to advance: %s", ctx.Err())

}

}

// Done returns a channel that is closed when all timers and ticks complete.

func (w AdvanceWaiter) Done() <-chan struct{} {

return w.ch

}

// Advance moves the clock forward by d, triggering any timers or tickers. The returned value can

// be used to wait for all timers and ticks to complete. Advance sets the clock forward before

// returning, and can only advance up to the next timer or tick event. It will fail the test if you

// attempt to advance beyond.

//

// If you need to advance exactly to the next event, and don't know or don't wish to calculate it,

// consider AdvanceNext().

func (m *Mock) Advance(d time.Duration) AdvanceWaiter {

m.tb.Helper()

w := AdvanceWaiter{tb: m.tb, ch: make(chan struct{})}

m.mu.Lock()

if !m.testOver {

m.logger.Logf("Mock Clock - Advance(%s)", d)

}

fin := m.cur.Add(d)

// nextTime.IsZero implies no events scheduled.

if m.nextTime.IsZero() || fin.Before(m.nextTime) {

m.cur = fin

m.mu.Unlock()

close(w.ch)

return w

}

if fin.After(m.nextTime) {

m.tb.Errorf("cannot advance %s which is beyond next timer/ticker event in %s",

d.String(), m.nextTime.Sub(m.cur))

m.mu.Unlock()

close(w.ch)

return w

}

m.cur = m.nextTime

go m.advanceLocked(w)

return w

}

func (m *Mock) advanceLocked(w AdvanceWaiter) {

defer close(w.ch)

wg := sync.WaitGroup{}

for i := range m.nextEvents {

e := m.nextEvents[i]

t := m.cur

wg.Add(1)

go func() {

e.fire(t)

wg.Done()

}()

}

// release the lock and let the events resolve. This allows them to call back into the

// Mock to query the time or set new timers. Each event should remove or reschedule

// itself from nextEvents.

m.mu.Unlock()

wg.Wait()

}

// Set the time to t. If the time is after the current mocked time, then this is equivalent to

// Advance() with the difference. You may only Set the time earlier than the current time before

// starting tickers and timers (e.g. at the start of your test case).

func (m *Mock) Set(t time.Time) AdvanceWaiter {

m.tb.Helper()

w := AdvanceWaiter{tb: m.tb, ch: make(chan struct{})}

m.mu.Lock()

if !m.testOver {

m.logger.Logf("Mock Clock - Set(%s)", t)

}

if t.Before(m.cur) {

defer close(w.ch)

defer m.mu.Unlock()

// past

if !m.nextTime.IsZero() {

m.tb.Error("Set mock clock to the past after timers/tickers started")

}

m.cur = t

return w

}

// future

// nextTime.IsZero implies no events scheduled.

if m.nextTime.IsZero() || t.Before(m.nextTime) {

defer close(w.ch)

defer m.mu.Unlock()

m.cur = t

return w

}

if t.After(m.nextTime) {

defer close(w.ch)

defer m.mu.Unlock()

m.tb.Errorf("cannot Set time to %s which is beyond next timer/ticker event at %s",

t.String(), m.nextTime)

return w

}

m.cur = m.nextTime

go m.advanceLocked(w)

return w

}

// AdvanceNext advances the clock to the next timer or tick event. It fails the test if there are

// none scheduled. It returns the duration the clock was advanced and a waiter that can be used to

// wait for the timer/tick event(s) to finish.

func (m *Mock) AdvanceNext() (time.Duration, AdvanceWaiter) {

m.mu.Lock()

if !m.testOver {

m.logger.Logf("Mock Clock - AdvanceNext()")

}

m.tb.Helper()

w := AdvanceWaiter{tb: m.tb, ch: make(chan struct{})}

if m.nextTime.IsZero() {

defer close(w.ch)

defer m.mu.Unlock()

m.tb.Error("cannot AdvanceNext because there are no timers or tickers running")

return 0, w

}

d := m.nextTime.Sub(m.cur)

m.cur = m.nextTime

go m.advanceLocked(w)

return d, w

}

// Peek returns the duration until the next ticker or timer event and the value

// true, or, if there are no running tickers or timers, it returns zero and

// false.

func (m *Mock) Peek() (d time.Duration, ok bool) {

m.mu.Lock()

defer m.mu.Unlock()

if m.nextTime.IsZero() {

return 0, false

}

return m.nextTime.Sub(m.cur), true

}

// Trapper allows the creation of Traps

type Trapper struct {

// mock is the underlying Mock. This is a thin wrapper around Mock so that

// we can have our interface look like mClock.Trap().NewTimer("foo")

mock *Mock

}

func (t Trapper) NewTimer(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionNewTimer, tags)

}

func (t Trapper) AfterFunc(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionAfterFunc, tags)

}

func (t Trapper) TimerStop(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTimerStop, tags)

}

func (t Trapper) TimerReset(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTimerReset, tags)

}

func (t Trapper) TickerFunc(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTickerFunc, tags)

}

func (t Trapper) TickerFuncWait(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTickerFuncWait, tags)

}

func (t Trapper) NewTicker(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionNewTicker, tags)

}

func (t Trapper) TickerStop(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTickerStop, tags)

}

func (t Trapper) TickerReset(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionTickerReset, tags)

}

func (t Trapper) Now(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionNow, tags)

}

func (t Trapper) Since(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionSince, tags)

}

func (t Trapper) Until(tags ...string) *Trap {

return t.mock.newTrap(clockFunctionUntil, tags)

}

func (m *Mock) Trap() Trapper {

return Trapper{m}

}

func (m *Mock) newTrap(fn clockFunction, tags []string) *Trap {

m.mu.Lock()

defer m.mu.Unlock()

if !m.testOver {

m.logger.Logf("Mock Clock - Trap %s(..., %v)", fn, tags)

}

tr := &Trap{

fn: fn,

tags: tags,

mock: m,

calls: make(chan *apiCall),

done: make(chan struct{}),

}

m.traps = append(m.traps, tr)

return tr

}

// WithLogger replaces the default testing logger with a custom one.

//

// This can be used to discard log messages with:

//

// quartz.NewMock(t).WithLogger(quartz.NoOpLogger)

func (m *Mock) WithLogger(l Logger) *Mock {

m.mu.Lock()

defer m.mu.Unlock()

m.logger = l

return m

}

// NewMock creates a new Mock with the time set to midnight UTC on Jan 1, 2024.

// You may re-set the time earlier than this, but only before timers or tickers

// are created.

func NewMock(tb TestingT) *Mock {

cur, err := time.Parse(time.RFC3339, "2024-01-01T00:00:00Z")

if err != nil {

panic(err)

}

m := &Mock{

tb: tb,

logger: tb,

cur: cur,

}

tb.Cleanup(func() {

m.mu.Lock()

defer m.mu.Unlock()

m.testOver = true

m.logger.Logf("Mock Clock - test cleanup; will no longer log clock events")

})

return m

}

var _ Clock = &Mock{}

type mockTickerFunc struct {

ctx context.Context

d time.Duration

f func() error

nxt time.Time

mock *Mock

// cond is a condition Locked on the main Mock.mu

cond *sync.Cond

// inProgress is true when we are actively calling f

inProgress bool

// done is true when the ticker exits

done bool

// err holds the error when the ticker exits

err error

}

func (m *mockTickerFunc) next() time.Time {

return m.nxt

}

func (m *mockTickerFunc) fire(_ time.Time) {

m.mock.mu.Lock()

if m.done {

m.mock.mu.Unlock()

return

}

m.nxt = m.nxt.Add(m.d)

m.mock.recomputeNextLocked()

// we need this check to happen after we've computed the next tick,

// otherwise it will be immediately rescheduled.

if m.inProgress {

m.mock.mu.Unlock()

return

}

m.inProgress = true

m.mock.mu.Unlock()

err := m.f()

m.mock.mu.Lock()

defer m.mock.mu.Unlock()

m.inProgress = false

m.cond.Broadcast() // wake up anything waiting for f to finish

if err != nil {

m.exitLocked(err)

}

}

func (m *mockTickerFunc) exitLocked(err error) {

if m.done {

return

}

m.done = true

m.err = err

m.mock.removeEventLocked(m)

m.cond.Broadcast()

}

func (m *mockTickerFunc) waitForCtx() {

<-m.ctx.Done()

m.mock.mu.Lock()

defer m.mock.mu.Unlock()

for m.inProgress {

m.cond.Wait()

}

m.exitLocked(m.ctx.Err())

}

func (m *mockTickerFunc) Wait(tags ...string) error {

m.mock.mu.Lock()

defer m.mock.mu.Unlock()

c := newCall(clockFunctionTickerFuncWait, tags)

m.mock.matchCallLocked(c)

defer close(c.complete)

for !m.done {

m.cond.Wait()

}

return m.err

}

var _ Waiter = &mockTickerFunc{}

type clockFunction int

const (

clockFunctionNewTimer clockFunction = iota

clockFunctionAfterFunc

clockFunctionTimerStop

clockFunctionTimerReset

clockFunctionTickerFunc

clockFunctionTickerFuncWait

clockFunctionNewTicker

clockFunctionTickerReset

clockFunctionTickerStop

clockFunctionNow

clockFunctionSince

clockFunctionUntil

)

func (c clockFunction) String() string {

switch c {

case clockFunctionNewTimer:

return "NewTimer"

case clockFunctionAfterFunc:

return "AfterFunc"

case clockFunctionTimerStop:

return "Timer.Stop"

case clockFunctionTimerReset:

return "Timer.Reset"

case clockFunctionTickerFunc:

return "TickerFunc"

case clockFunctionTickerFuncWait:

return "TickerFunc.Wait"

case clockFunctionNewTicker:

return "NewTicker"

case clockFunctionTickerReset:

return "Ticker.Reset"

case clockFunctionTickerStop:

return "Ticker.Stop"

case clockFunctionNow:

return "Now"

case clockFunctionSince:

return "Since"

case clockFunctionUntil:

return "Until"

default:

return fmt.Sprintf("Unknown clockFunction(%d)", c)

}

}

type callArg func(c *apiCall)

// apiCall represents a single call to one of the Clock APIs.

type apiCall struct {

Time time.Time

Duration time.Duration

Tags []string

fn clockFunction

releases sync.WaitGroup

complete chan struct{}

}

func (a *apiCall) String() string {

switch a.fn {

case clockFunctionNewTimer:

return fmt.Sprintf("NewTimer(%s, %v)", a.Duration, a.Tags)

case clockFunctionAfterFunc:

return fmt.Sprintf("AfterFunc(%s, <fn>, %v)", a.Duration, a.Tags)

case clockFunctionTimerStop:

return fmt.Sprintf("Timer.Stop(%v)", a.Tags)

case clockFunctionTimerReset:

return fmt.Sprintf("Timer.Reset(%s, %v)", a.Duration, a.Tags)

case clockFunctionTickerFunc:

return fmt.Sprintf("TickerFunc(<ctx>, %s, <fn>, %s)", a.Duration, a.Tags)

case clockFunctionTickerFuncWait:

return fmt.Sprintf("TickerFunc.Wait(%v)", a.Tags)

case clockFunctionNewTicker:

return fmt.Sprintf("NewTicker(%s, %v)", a.Duration, a.Tags)

case clockFunctionTickerReset:

return fmt.Sprintf("Ticker.Reset(%s, %v)", a.Duration, a.Tags)

case clockFunctionTickerStop:

return fmt.Sprintf("Ticker.Stop(%v)", a.Tags)

case clockFunctionNow:

return fmt.Sprintf("Now(%v)", a.Tags)

case clockFunctionSince:

return fmt.Sprintf("Since(%s, %v)", a.Time, a.Tags)

case clockFunctionUntil:

return fmt.Sprintf("Until(%s, %v)", a.Time, a.Tags)

default:

return fmt.Sprintf("Unknown clockFunction(%d)", a.fn)

}

}

// Call represents an apiCall that has been trapped.

type Call struct {

Time time.Time

Duration time.Duration

Tags []string

tb TestingT

apiCall *apiCall

trap *Trap

}

// Release the call and wait for it to complete. If the provided context expires before the call completes, it returns

// an error.

//

// IMPORTANT: If a call is trapped by more than one trap, they all must release the call before it can complete, and

// they must do so from different goroutines.

func (c *Call) Release(ctx context.Context) error {

c.apiCall.releases.Done()

select {

case <-ctx.Done():

return fmt.Errorf("timed out waiting for release; did more than one trap capture the call?: %w", ctx.Err())

case <-c.apiCall.complete:

// OK

}

c.trap.callReleased()

return nil

}

// MustRelease releases the call and waits for it to complete. If the provided context expires before the call

// completes, it fails the test.

//

// IMPORTANT: If a call is trapped by more than one trap, they all must release the call before it can complete, and

// they must do so from different goroutines.

func (c *Call) MustRelease(ctx context.Context) {

if err := c.Release(ctx); err != nil {

c.tb.Helper()

c.tb.Fatal(err.Error())

}

}

func withTime(t time.Time) callArg {

return func(c *apiCall) {

c.Time = t

}

}

func withDuration(d time.Duration) callArg {

return func(c *apiCall) {

c.Duration = d

}

}

func newCall(fn clockFunction, tags []string, args ...callArg) *apiCall {

c := &apiCall{

fn: fn,

Tags: tags,

complete: make(chan struct{}),

}

for _, a := range args {

a(c)

}

return c

}

type Trap struct {

fn clockFunction

tags []string

mock *Mock

calls chan *apiCall

done chan struct{}

// mu protects the unreleasedCalls count

mu sync.Mutex

unreleasedCalls int

}

func (t *Trap) String() string {

return fmt.Sprintf("Trap %s(..., %v)", t.fn.String(), t.tags)

}

func (t *Trap) catch(c *apiCall) {

select {

case t.calls <- c:

case <-t.done:

c.releases.Done()

}

}

func (t *Trap) matches(c *apiCall) bool {

if t.fn != c.fn {

return false

}

for _, tag := range t.tags {

if !slices.Contains(c.Tags, tag) {

return false

}

}

return true

}

func (t *Trap) Close() {

t.mock.mu.Lock()

defer t.mock.mu.Unlock()

select {

case <-t.done:

t.mock.tb.Logf("%s already Closed()", t)

return // already closed

default:

}

if t.unreleasedCalls != 0 {

t.mock.tb.Helper()

t.mock.tb.Errorf("%s Closed() with %d unreleased calls", t, t.unreleasedCalls)

}

for i, tr := range t.mock.traps {

if t == tr {

t.mock.traps = append(t.mock.traps[:i], t.mock.traps[i+1:]...)

}

}

close(t.done)

}

func (t *Trap) callReleased() {

t.mu.Lock()

defer t.mu.Unlock()

t.unreleasedCalls--

}

var ErrTrapClosed = errors.New("trap closed")

func (t *Trap) Wait(ctx context.Context) (*Call, error) {

select {

case <-ctx.Done():

return nil, ctx.Err()

case <-t.done:

return nil, ErrTrapClosed

case a := <-t.calls:

c := &Call{

Time: a.Time,

Duration: a.Duration,

Tags: a.Tags,

apiCall: a,

trap: t,

tb: t.mock.tb,

}

t.mu.Lock()

defer t.mu.Unlock()

t.unreleasedCalls++

return c, nil

}

}

// MustWait calls Wait() and then if there is an error, immediately fails the

// test via tb.Fatalf()

func (t *Trap) MustWait(ctx context.Context) *Call {

t.mock.tb.Helper()

c, err := t.Wait(ctx)

if err != nil {

t.mock.tb.Fatalf("context expired while waiting for %s: %s", t, err.Error())

}

return c

}

type Logger interface {

Log(args ...any)

Logf(format string, args ...any)

}

// NoOpLogger is a Logger that discards all log messages.

var NoOpLogger Logger = noOpLogger{}

type noOpLogger struct{}

func (noOpLogger) Log(args ...any) {}

func (noOpLogger) Logf(format string, args ...any) {}