grafana/pkg/services/ngalert/notifier/images_test.go

package notifier

import (
	"context"
	"io"
	"os"
	"path/filepath"
	"testing"
	"time"

	alertingImages "github.com/grafana/alerting/images"
	alertingModels "github.com/grafana/alerting/models"
	alertingNotify "github.com/grafana/alerting/notify"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/ngalert/store"
	"github.com/prometheus/common/model"
	"github.com/stretchr/testify/require"
)

func TestGetImage(t *testing.T) {
	fakeImageStore := store.NewFakeImageStore(t)
	store := newImageProvider(fakeImageStore, log.NewNopLogger())

	t.Run("queries by token when it gets a token", func(tt *testing.T) {
		img := models.Image{
			Token: "test",
			URL:   "http://localhost:1234",
			Path:  "test.png",
		}
		err := fakeImageStore.SaveImage(context.Background(), &img)
		require.NoError(tt, err)

		// nolint:staticcheck
		savedImg, err := store.GetImage(context.Background(), "token://"+img.Token)
		require.NoError(tt, err)
		require.Equal(tt, savedImg.Token, img.Token)
		require.Equal(tt, savedImg.URL, img.URL)
		require.Equal(tt, savedImg.Path, img.Path)
	})

	t.Run("queries by URL when it gets a URL", func(tt *testing.T) {
		img := models.Image{
			Token: "test",
			Path:  "test.png",
			URL:   "https://test.com/test.png",
		}
		err := fakeImageStore.SaveImage(context.Background(), &img)
		require.NoError(tt, err)

		// nolint:staticcheck
		savedImg, err := store.GetImage(context.Background(), img.URL)
		require.NoError(tt, err)
		require.Equal(tt, savedImg.Token, img.Token)
		require.Equal(tt, savedImg.URL, img.URL)
		require.Equal(tt, savedImg.Path, img.Path)
	})
}

func TestGetImageURL(t *testing.T) {
	var (
		imageWithoutURL = models.Image{
			Token:     "test-no-url",
			CreatedAt: time.Now().UTC(),
			ExpiresAt: time.Now().UTC().Add(24 * time.Hour),
		}
		testImage = models.Image{
			Token:     "test",
			URL:       "https://test.com",
			CreatedAt: time.Now().UTC(),
			ExpiresAt: time.Now().UTC().Add(24 * time.Hour),
		}
	)

	fakeImageStore := store.NewFakeImageStore(t, &imageWithoutURL, &testImage)
	store := newImageProvider(fakeImageStore, log.NewNopLogger())

	tests := []struct {
		name   string
		uri    string
		expURL string
		expErr error
	}{
		{
			"URL does not exist",
			"https://invalid.com/test",
			"",
			alertingImages.ErrImageNotFound,
		}, {
			"existing URL",
			testImage.URL,
			testImage.URL,
			nil,
		}, {
			"token does not exist",
			"token://invalid",
			"",
			alertingImages.ErrImageNotFound,
		}, {
			"existing token",
			"token://" + testImage.Token,
			testImage.URL,
			nil,
		}, {
			"image has no URL",
			"token://" + imageWithoutURL.Token,
			"",
			alertingImages.ErrImagesNoURL,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(tt *testing.T) {
			alert := alertingNotify.Alert{
				Alert: model.Alert{
					Annotations: model.LabelSet{alertingModels.ImageTokenAnnotation: model.LabelValue(test.uri)},
				},
			}
			url, err := store.GetImageURL(context.Background(), &alert)
			require.ErrorIs(tt, err, test.expErr)
			require.Equal(tt, test.expURL, url)
		})
	}
}

func TestGetRawImage(t *testing.T) {
	var (
		testBytes        = []byte("some test bytes")
		testPath         = generateTestFile(t, testBytes)
		imageWithoutPath = models.Image{
			Token:     "test-no-path",
			URL:       "https://test-no-path.com",
			CreatedAt: time.Now().UTC(),
			ExpiresAt: time.Now().UTC().Add(24 * time.Hour),
		}
		testImage = models.Image{
			Token:     "test",
			URL:       "https://test.com",
			Path:      testPath,
			CreatedAt: time.Now().UTC(),
			ExpiresAt: time.Now().UTC().Add(24 * time.Hour),
		}
	)

	fakeImageStore := store.NewFakeImageStore(t, &imageWithoutPath, &testImage)
	store := newImageProvider(fakeImageStore, log.NewNopLogger())

	tests := []struct {
		name        string
		uri         string
		expFilename string
		expBytes    []byte
		expErr      error
	}{
		{
			"URL does not exist",
			"https://invalid.com/test",
			"",
			nil,
			alertingImages.ErrImageNotFound,
		}, {
			"existing URL",
			testImage.URL,
			filepath.Base(testPath),
			testBytes,
			nil,
		}, {
			"token does not exist",
			"token://invalid",
			"",
			nil,
			alertingImages.ErrImageNotFound,
		}, {
			"existing token",
			"token://" + testImage.Token,
			filepath.Base(testPath),
			testBytes,
			nil,
		}, {
			"image has no path",
			"token://" + imageWithoutPath.Token,
			"",
			nil,
			alertingImages.ErrImagesNoPath,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(tt *testing.T) {
			alert := alertingNotify.Alert{
				Alert: model.Alert{
					Annotations: model.LabelSet{alertingModels.ImageTokenAnnotation: model.LabelValue(test.uri)},
				},
			}
			readCloser, filename, err := store.GetRawImage(context.Background(), &alert)
			require.ErrorIs(tt, err, test.expErr)
			require.Equal(tt, test.expFilename, filename)

			if test.expBytes != nil {
				b, err := io.ReadAll(readCloser)
				require.NoError(tt, err)
				require.Equal(tt, test.expBytes, b)
				require.NoError(t, readCloser.Close())
			}
		})
	}
}

func generateTestFile(t *testing.T, b []byte) string {
	t.Helper()
	f, err := os.CreateTemp("/tmp", "image")
	require.NoError(t, err)
	defer func(f *os.File) {
		_ = f.Close()
	}(f)

	t.Cleanup(func() {
		require.NoError(t, os.RemoveAll(f.Name()))
	})

	_, err = f.Write(b)
	require.NoError(t, err)

	return f.Name()
}