@@ -3,6 +3,7 @@ package notifications_test
33import (
44"context"
55"strconv"
6+ "sync"
67"testing"
78"time"
89
@@ -317,10 +318,12 @@ func TestInflightDispatchesMetric(t *testing.T) {
317318})
318319
319320handler := & fakeHandler {}
320- // Delayer will delay all dispatches by 2x fetch intervals to ensure we catch the requests inflight.
321- delayer := newDelayingHandler (cfg .FetchInterval .Value ()* 2 ,handler )
321+ const msgCount = 2
322+
323+ // Barrier handler will wait until all notification messages are in-flight.
324+ barrier := newBarrierHandler (msgCount ,handler )
322325mgr .WithHandlers (map [database.NotificationMethod ]notifications.Handler {
323- method :delayer ,
326+ method :barrier ,
324327})
325328
326329enq ,err := notifications .NewStoreEnqueuer (cfg ,api .Database ,defaultHelpers (),api .Logger .Named ("enqueuer" ),quartz .NewReal ())
@@ -329,7 +332,6 @@ func TestInflightDispatchesMetric(t *testing.T) {
329332user := createSampleUser (t ,api .Database )
330333
331334// WHEN: notifications are enqueued which will succeed (and be delayed during dispatch)
332- const msgCount = 2
333335for i := 0 ;i < msgCount ;i ++ {
334336_ ,err = enq .Enqueue (ctx ,user .ID ,template ,map [string ]string {"type" :"success" ,"i" :strconv .Itoa (i )},"test" )
335337require .NoError (t ,err )
@@ -343,6 +345,10 @@ func TestInflightDispatchesMetric(t *testing.T) {
343345return promtest .ToFloat64 (metrics .InflightDispatches .WithLabelValues (string (method ),template .String ()))== msgCount
344346},testutil .WaitShort ,testutil .IntervalFast )
345347
348+ for i := 0 ;i < msgCount ;i ++ {
349+ barrier .wg .Done ()
350+ }
351+
346352// Wait until the handler has dispatched the given notifications.
347353require .Eventually (t ,func ()bool {
348354handler .mu .RLock ()
@@ -493,27 +499,30 @@ func (u *updateSignallingInterceptor) BulkMarkNotificationMessagesFailed(ctx con
493499return u .Store .BulkMarkNotificationMessagesFailed (ctx ,arg )
494500}
495501
496- type delayingHandler struct {
502+ type barrierHandler struct {
497503h notifications.Handler
498504
499- delay time. Duration
505+ wg * sync. WaitGroup
500506}
501507
502- func newDelayingHandler (delay time.Duration ,handler notifications.Handler )* delayingHandler {
503- return & delayingHandler {
504- delay :delay ,
505- h :handler ,
508+ func newBarrierHandler (total int ,handler notifications.Handler )* barrierHandler {
509+ var wg sync.WaitGroup
510+ wg .Add (total )
511+
512+ return & barrierHandler {
513+ h :handler ,
514+ wg :& wg ,
506515}
507516}
508517
509- func (d * delayingHandler )Dispatcher (payload types.MessagePayload ,title ,body string ) (dispatch.DeliveryFunc ,error ) {
510- deliverFn ,err := d .h .Dispatcher (payload ,title ,body )
518+ func (bh * barrierHandler )Dispatcher (payload types.MessagePayload ,title ,body string ) (dispatch.DeliveryFunc ,error ) {
519+ deliverFn ,err := bh .h .Dispatcher (payload ,title ,body )
511520if err != nil {
512521return nil ,err
513522}
514523
515524return func (ctx context.Context ,msgID uuid.UUID ) (retryable bool ,err error ) {
516- time . Sleep ( d . delay )
525+ bh . wg . Wait ( )
517526
518527return deliverFn (ctx ,msgID )
519528},nil