Alerting: Add ErrImagesDone to return from withStoredImages (#51098)

pkg/services/ngalert/notifier/channels/slack.go

@@ -323,6 +323,7 @@ func (sn *SlackNotifier) buildSlackMessage(ctx context.Context, alrts []*types.A
 		func(index int, image *ngmodels.Image) error {
 			if image != nil {
 				req.Attachments[0].ImageURL = image.URL
+				return ErrImagesDone
 			}
 			return nil
 		},

pkg/services/ngalert/notifier/channels/util.go

@@ -37,7 +37,12 @@ const (
 
 var (
 	// Provides current time. Can be overwritten in tests.
-	timeNow              = time.Now
+	timeNow = time.Now
+
+	// ErrImagesDone is used to stop iteration of subsequent images. It should be
+	// returned from forEachFunc when either the intended image has been found or
+	// the maximum number of images has been iterated.
+	ErrImagesDone        = errors.New("images done")
 	ErrImagesUnavailable = errors.New("alert screenshots are unavailable")
 )
 
@@ -53,6 +58,11 @@ func withStoredImages(ctx context.Context, l log.Logger, imageStore ImageStore,
 	for i := range alerts {
 		err := withStoredImage(ctx, l, imageStore, forEachFunc, i, alerts...)
 		if err != nil {
+			// Stop iteration as forEachFunc has found the intended image or
+			// iterated the maximum number of images
+			if errors.Is(err, ErrImagesDone) {
+				return nil
+			}
 			return err
 		}
 	}

pkg/services/ngalert/notifier/channels/util_test.go

@@ -0,0 +1,63 @@
+package channels
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/prometheus/alertmanager/types"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+)
+
+func TestWithStoredImages(t *testing.T) {
+	ctx := context.Background()
+	alerts := []*types.Alert{{
+		Alert: model.Alert{
+			Annotations: model.LabelSet{
+				models.ScreenshotTokenAnnotation: "test-image-1",
+			},
+		},
+	}, {
+		Alert: model.Alert{
+			Annotations: model.LabelSet{
+				models.ScreenshotTokenAnnotation: "test-image-2",
+			},
+		},
+	}}
+	imageStore := &fakeImageStore{Images: []*models.Image{{
+		Token:     "test-image-1",
+		URL:       "https://www.example.com/test-image-1.jpg",
+		CreatedAt: time.Now().UTC(),
+	}, {
+		Token:     "test-image-2",
+		URL:       "https://www.example.com/test-image-2.jpg",
+		CreatedAt: time.Now().UTC(),
+	}}}
+
+	var (
+		err error
+		i   int
+	)
+
+	// should iterate all images
+	err = withStoredImages(ctx, log.New(ctx), imageStore, func(index int, image *models.Image) error {
+		i += 1
+		return nil
+	}, alerts...)
+	require.NoError(t, err)
+	assert.Equal(t, 2, i)
+
+	// should iterate just the first image
+	i = 0
+	err = withStoredImages(ctx, log.New(ctx), imageStore, func(index int, image *models.Image) error {
+		i += 1
+		return ErrImagesDone
+	}, alerts...)
+	require.NoError(t, err)
+	assert.Equal(t, 1, i)
+}