From 62c2b1ec78e259ed9f4a7177a1dd665a76083bfa Mon Sep 17 00:00:00 2001
From: George Robinson <george.robinson@grafana.com>
Date: Mon, 20 Jun 2022 10:56:28 +0100
Subject: [PATCH] Alerting: Add ErrImagesDone to return from withStoredImages
 (#51098)

---
 .../ngalert/notifier/channels/slack.go        |  1 +
 .../notifier/channels/{utils.go => util.go}   | 12 +++-
 .../ngalert/notifier/channels/util_test.go    | 63 +++++++++++++++++++
 3 files changed, 75 insertions(+), 1 deletion(-)
 rename pkg/services/ngalert/notifier/channels/{utils.go => util.go} (93%)
 create mode 100644 pkg/services/ngalert/notifier/channels/util_test.go

diff --git a/pkg/services/ngalert/notifier/channels/slack.go b/pkg/services/ngalert/notifier/channels/slack.go
index 1dfefe585ec..282e4f73857 100644
--- a/pkg/services/ngalert/notifier/channels/slack.go
+++ b/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
 		},
diff --git a/pkg/services/ngalert/notifier/channels/utils.go b/pkg/services/ngalert/notifier/channels/util.go
similarity index 93%
rename from pkg/services/ngalert/notifier/channels/utils.go
rename to pkg/services/ngalert/notifier/channels/util.go
index ba1256048b9..ab9ef83bc18 100644
--- a/pkg/services/ngalert/notifier/channels/utils.go
+++ b/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
 		}
 	}
diff --git a/pkg/services/ngalert/notifier/channels/util_test.go b/pkg/services/ngalert/notifier/channels/util_test.go
new file mode 100644
index 00000000000..beec5dc848f
--- /dev/null
+++ b/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)
+}