grafana/pkg/services/annotations/annotationsimpl/composite_store_test.go

package annotationsimpl

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/annotations"
	"github.com/grafana/grafana/pkg/services/annotations/accesscontrol"
)

var (
	errGet     = errors.New("get error")
	errGetTags = errors.New("get tags error")
)

func TestCompositeStore(t *testing.T) {
	t.Run("should handle panic", func(t *testing.T) {
		r1 := newFakeReader()
		getPanic := func(context.Context, annotations.ItemQuery, *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error) {
			panic("ohno")
		}
		r2 := newFakeReader(withGetFn(getPanic))
		store := &CompositeStore{
			log.NewNopLogger(),
			[]readStore{r1, r2},
		}

		_, err := store.Get(context.Background(), annotations.ItemQuery{}, nil)
		require.Error(t, err)
		require.Contains(t, err.Error(), "concurrent job panic")
	})

	t.Run("should return first error", func(t *testing.T) {
		err1 := errors.New("error 1")
		r1 := newFakeReader(withError(err1))
		err2 := errors.New("error 2")
		r2 := newFakeReader(withError(err2), withWait(10*time.Millisecond))

		store := &CompositeStore{
			log.NewNopLogger(),
			[]readStore{r1, r2},
		}

		tc := []struct {
			f   func() (any, error)
			err error
		}{
			{
				f:   func() (any, error) { return store.Get(context.Background(), annotations.ItemQuery{}, nil) },
				err: errGet,
			},
			{
				f:   func() (any, error) { return store.GetTags(context.Background(), annotations.TagsQuery{}) },
				err: errGetTags,
			},
		}

		for _, tt := range tc {
			_, err := tt.f()
			require.Error(t, err)
			require.ErrorIs(t, err, err1)
			require.NotErrorIs(t, err, err2)
		}
	})

	t.Run("should combine and sort results from Get", func(t *testing.T) {
		items1 := []*annotations.ItemDTO{
			{TimeEnd: 1, Time: 2},
			{TimeEnd: 2, Time: 1},
		}
		r1 := newFakeReader(withItems(items1))

		items2 := []*annotations.ItemDTO{
			{TimeEnd: 1, Time: 1},
			{TimeEnd: 1, Time: 3},
		}
		r2 := newFakeReader(withItems(items2))

		store := &CompositeStore{
			log.NewNopLogger(),
			[]readStore{r1, r2},
		}

		expected := []*annotations.ItemDTO{
			{TimeEnd: 2, Time: 1},
			{TimeEnd: 1, Time: 3},
			{TimeEnd: 1, Time: 2},
			{TimeEnd: 1, Time: 1},
		}

		items, _ := store.Get(context.Background(), annotations.ItemQuery{}, nil)
		require.Equal(t, expected, items)
	})

	t.Run("should combine and sort results from GetTags", func(t *testing.T) {
		tags1 := []*annotations.TagsDTO{
			{Tag: "key1:val1"},
			{Tag: "key2:val1"},
		}
		r1 := newFakeReader(withTags(tags1))

		tags2 := []*annotations.TagsDTO{
			{Tag: "key1:val2"},
			{Tag: "key2:val2"},
		}
		r2 := newFakeReader(withTags(tags2))

		store := &CompositeStore{
			log.NewNopLogger(),
			[]readStore{r1, r2},
		}

		expected := []*annotations.TagsDTO{
			{Tag: "key1:val1"},
			{Tag: "key1:val2"},
			{Tag: "key2:val1"},
			{Tag: "key2:val2"},
		}

		res, _ := store.GetTags(context.Background(), annotations.TagsQuery{})
		require.Equal(t, expected, res.Tags)
	})

	// Check if reader is not modifying query since it might cause a race condition in case of composite store
	t.Run("should not modify query", func(t *testing.T) {
		getFn1 := func(ctx context.Context, query annotations.ItemQuery, resources *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error) {
			query.From = 1
			return []*annotations.ItemDTO{}, nil
		}
		getFn2 := func(ctx context.Context, query annotations.ItemQuery, resources *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error) {
			return []*annotations.ItemDTO{}, nil
		}

		r1 := newFakeReader(withGetFn(getFn1))
		r2 := newFakeReader(withGetFn(getFn2))

		store := &CompositeStore{
			log.NewNopLogger(),
			[]readStore{r1, r2},
		}

		query := annotations.ItemQuery{}
		_, err := store.Get(context.Background(), query, nil)
		require.NoError(t, err)
		require.Equal(t, int64(0), query.From)
	})
}

type fakeReader struct {
	items    []*annotations.ItemDTO
	tagRes   annotations.FindTagsResult
	getFn    func(context.Context, annotations.ItemQuery, *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error)
	getTagFn func(context.Context, annotations.TagsQuery) (annotations.FindTagsResult, error)
	wait     time.Duration
	err      error
}

func (f *fakeReader) Type() string {
	return "fake"
}

func (f *fakeReader) Get(ctx context.Context, query annotations.ItemQuery, accessResources *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error) {
	if f.getFn != nil {
		return f.getFn(ctx, query, accessResources)
	}

	if f.wait > 0 {
		time.Sleep(f.wait)
	}

	if f.err != nil {
		err := fmt.Errorf("%w: %w", errGet, f.err)
		return nil, err
	}

	return f.items, nil
}

func (f *fakeReader) GetTags(ctx context.Context, query annotations.TagsQuery) (annotations.FindTagsResult, error) {
	if f.getTagFn != nil {
		return f.getTagFn(ctx, query)
	}

	if f.wait > 0 {
		time.Sleep(f.wait)
	}

	if f.err != nil {
		err := fmt.Errorf("%w: %w", errGetTags, f.err)
		return annotations.FindTagsResult{}, err
	}

	return f.tagRes, nil
}

func withWait(wait time.Duration) func(*fakeReader) {
	return func(f *fakeReader) {
		f.wait = wait
	}
}

func withError(err error) func(*fakeReader) {
	return func(f *fakeReader) {
		f.err = err
	}
}

func withItems(items []*annotations.ItemDTO) func(*fakeReader) {
	return func(f *fakeReader) {
		f.items = items
	}
}

func withTags(tags []*annotations.TagsDTO) func(*fakeReader) {
	return func(f *fakeReader) {
		f.tagRes = annotations.FindTagsResult{Tags: tags}
	}
}

func withGetFn(fn func(context.Context, annotations.ItemQuery, *accesscontrol.AccessResources) ([]*annotations.ItemDTO, error)) func(*fakeReader) {
	return func(f *fakeReader) {
		f.getFn = fn
	}
}

func newFakeReader(opts ...func(*fakeReader)) *fakeReader {
	f := &fakeReader{}
	for _, opt := range opts {
		opt(f)
	}
	return f
}