fix: fix mock TickerFunc to act like real

mock.go

@@ -466,6 +466,8 @@ type mockTickerFunc struct {
 
 	// 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
@@ -479,15 +481,18 @@ func (m *mockTickerFunc) next() time.Time {
 func (m *mockTickerFunc) fire(_ time.Time) {
 	m.mock.mu.Lock()
 	defer m.mock.mu.Unlock()
-	if m.done {
+	if m.done || m.inProgress {
 		return
 	}
 	m.nxt = m.nxt.Add(m.d)
 	m.mock.recomputeNextLocked()
 
+	m.inProgress = true
 	m.mock.mu.Unlock()
 	err := m.f()
 	m.mock.mu.Lock()
+	m.inProgress = false
+	m.cond.Broadcast() // wake up anything waiting for f to finish
 	if err != nil {
 		m.exitLocked(err)
 	}
@@ -507,6 +512,9 @@ 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())
 }
 

mock_test.go

@@ -2,6 +2,7 @@ package quartz_test
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -210,3 +211,107 @@ func TestPeek(t *testing.T) {
 		t.Fatal("expected Peek() to return false")
 	}
 }
+
+// TestTickerFunc_ContextDoneDuringTick tests that TickerFunc.Wait() can't return while the tick
+// function callback is in progress.
+func TestTickerFunc_ContextDoneDuringTick(t *testing.T) {
+	t.Parallel()
+	testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer testCancel()
+	mClock := quartz.NewMock(t)
+
+	tickStart := make(chan struct{})
+	tickDone := make(chan struct{})
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+	tkr := mClock.TickerFunc(ctx, time.Second, func() error {
+		close(tickStart)
+		select {
+		case <-tickDone:
+		case <-testCtx.Done():
+			t.Error("timeout waiting for tickDone")
+		}
+		return nil
+	})
+	w := mClock.Advance(time.Second)
+	select {
+	case <-tickStart:
+		// OK
+	case <-testCtx.Done():
+		t.Fatal("timeout waiting for tickStart")
+	}
+	waitErr := make(chan error, 1)
+	go func() {
+		waitErr <- tkr.Wait()
+	}()
+	cancel()
+	select {
+	case <-waitErr:
+		t.Fatal("wait should not return while tick callback in progress")
+	case <-time.After(time.Millisecond * 100):
+		// OK
+	}
+	close(tickDone)
+	select {
+	case err := <-waitErr:
+		if !errors.Is(err, context.Canceled) {
+			t.Fatal("expected context.Canceled error")
+		}
+	case <-testCtx.Done():
+		t.Fatal("timed out waiting for wait to finish")
+	}
+	w.MustWait(testCtx)
+}
+
+// TestTickerFunc_LongCallback tests that we don't call the ticker func a second time while the
+// first is still executing.
+func TestTickerFunc_LongCallback(t *testing.T) {
+	t.Parallel()
+	testCtx, testCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer testCancel()
+	mClock := quartz.NewMock(t)
+
+	expectedErr := errors.New("callback error")
+	tickStart := make(chan struct{})
+	tickDone := make(chan struct{})
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+	tkr := mClock.TickerFunc(ctx, time.Second, func() error {
+		close(tickStart)
+		select {
+		case <-tickDone:
+		case <-testCtx.Done():
+			t.Error("timeout waiting for tickDone")
+		}
+		return expectedErr
+	})
+	w := mClock.Advance(time.Second)
+	select {
+	case <-tickStart:
+		// OK
+	case <-testCtx.Done():
+		t.Fatal("timeout waiting for tickStart")
+	}
+	// second tick completes immediately, since it doesn't actually call the
+	// ticker function.
+	mClock.Advance(time.Second).MustWait(testCtx)
+
+	waitErr := make(chan error, 1)
+	go func() {
+		waitErr <- tkr.Wait()
+	}()
+	cancel()
+	close(tickDone)
+
+	select {
+	case err := <-waitErr:
+		// we should get the function error, not the context error, since context was canceled while
+		// we were calling the function, and it returned an error.
+		if !errors.Is(err, expectedErr) {
+			t.Fatalf("wrong error: %s", err)
+		}
+	case <-testCtx.Done():
+		t.Fatal("timed out waiting for wait to finish")
+	}
+	w.MustWait(testCtx)
+}