mirror of
https://github.com/grafana/grafana.git
synced 2026-08-18 17:15:08 -05:00
Alerting: Add a general screenshot service and alerting-specific image service. (#49293)
This commit adds a pkg/services/screenshot package for taking and uploading screenshots of Grafana dashboards. It supports taking screenshots of both dashboards and individual panels within a dashboard, using the rendering service. The screenshot package has the following services, most of which can be composed: BrowserScreenshotService (Takes screenshots with headless Chrome) CachableScreenshotService (Caches screenshots taken with another service such as BrowserScreenshotService) NoopScreenshotService (A no-op screenshot service for tests) SingleFlightScreenshotService (Prevents duplicate screenshots when taking screenshots of the same dashboard or panel in parallel) ScreenshotUnavailableService (A screenshot service that returns ErrScreenshotsUnavailable) UploadingScreenshotService (A screenshot service that uploads taken screenshots) The screenshot package does not support wire dependency injection yet. ngalert constructs its own version of the service. See https://github.com/grafana/grafana/issues/49296 This PR also adds an ImageScreenshotService to ngAlert. This is used to take screenshots with a screenshotservice and then store their location reference for use by alert instances and notifiers.
This commit is contained in:
Joe Blubaugh
@@ -0,0 +1,50 @@
|
||||
// Code generated by MockGen. DO NOT EDIT.
|
||||
// Source: github.com/grafana/grafana/pkg/services/screenshot (interfaces: ScreenshotService)
|
||||
|
||||
// Package screenshot is a generated GoMock package.
|
||||
package screenshot
|
||||
|
||||
import (
|
||||
context "context"
|
||||
reflect "reflect"
|
||||
|
||||
gomock "github.com/golang/mock/gomock"
|
||||
)
|
||||
|
||||
// MockScreenshotService is a mock of ScreenshotService interface.
|
||||
type MockScreenshotService struct {
|
||||
ctrl *gomock.Controller
|
||||
recorder *MockScreenshotServiceMockRecorder
|
||||
}
|
||||
|
||||
// MockScreenshotServiceMockRecorder is the mock recorder for MockScreenshotService.
|
||||
type MockScreenshotServiceMockRecorder struct {
|
||||
mock *MockScreenshotService
|
||||
}
|
||||
|
||||
// NewMockScreenshotService creates a new mock instance.
|
||||
func NewMockScreenshotService(ctrl *gomock.Controller) *MockScreenshotService {
|
||||
mock := &MockScreenshotService{ctrl: ctrl}
|
||||
mock.recorder = &MockScreenshotServiceMockRecorder{mock}
|
||||
return mock
|
||||
}
|
||||
|
||||
// EXPECT returns an object that allows the caller to indicate expected use.
|
||||
func (m *MockScreenshotService) EXPECT() *MockScreenshotServiceMockRecorder {
|
||||
return m.recorder
|
||||
}
|
||||
|
||||
// Take mocks base method.
|
||||
func (m *MockScreenshotService) Take(arg0 context.Context, arg1 ScreenshotOptions) (*Screenshot, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "Take", arg0, arg1)
|
||||
ret0, _ := ret[0].(*Screenshot)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// Take indicates an expected call of Take.
|
||||
func (mr *MockScreenshotServiceMockRecorder) Take(arg0, arg1 interface{}) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Take", reflect.TypeOf((*MockScreenshotService)(nil).Take), arg0, arg1)
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
package screenshot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
gocache "github.com/patrickmn/go-cache"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/promauto"
|
||||
"golang.org/x/sync/singleflight"
|
||||
|
||||
"github.com/grafana/grafana/pkg/components/imguploader"
|
||||
"github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/services/dashboards"
|
||||
"github.com/grafana/grafana/pkg/services/rendering"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
)
|
||||
|
||||
const (
|
||||
namespace = "grafana"
|
||||
subsystem = "screenshot"
|
||||
)
|
||||
|
||||
var (
|
||||
DefaultTheme = models.ThemeDark
|
||||
DefaultTimeout = 15 * time.Second
|
||||
DefaultHeight = 500
|
||||
DefaultWidth = 1000
|
||||
)
|
||||
|
||||
var (
|
||||
ErrScreenshotsUnavailable = errors.New("screenshots unavailable")
|
||||
)
|
||||
|
||||
// Screenshot represents a screenshot of a dashboard in Grafana.
|
||||
//
|
||||
// A screenshot can have a Path and an URL if the screenshot is stored on disk
|
||||
// and uploaded to a cloud storage service or made accessible via the Grafana
|
||||
// HTTP server.
|
||||
type Screenshot struct {
|
||||
Path string
|
||||
URL string
|
||||
}
|
||||
|
||||
// ScreenshotOptions are the options for taking a screenshot.
|
||||
type ScreenshotOptions struct {
|
||||
DashboardUID string
|
||||
PanelID int64
|
||||
Width int
|
||||
Height int
|
||||
Theme models.Theme
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
// SetDefaults sets default values for missing or invalid options.
|
||||
func (s ScreenshotOptions) SetDefaults() ScreenshotOptions {
|
||||
if s.Width <= 0 {
|
||||
s.Width = DefaultWidth
|
||||
}
|
||||
if s.Height <= 0 {
|
||||
s.Height = DefaultHeight
|
||||
}
|
||||
switch s.Theme {
|
||||
case models.ThemeDark, models.ThemeLight:
|
||||
default:
|
||||
s.Theme = DefaultTheme
|
||||
}
|
||||
if s.Timeout <= 0 {
|
||||
s.Timeout = DefaultTimeout
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// ScreenshotService is an interface for taking screenshots.
|
||||
//go:generate mockgen -destination=mock.go -package=screenshot github.com/grafana/grafana/pkg/services/screenshot ScreenshotService
|
||||
type ScreenshotService interface {
|
||||
Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error)
|
||||
}
|
||||
|
||||
// BrowserScreenshotService takes screenshots using a headless browser.
|
||||
type BrowserScreenshotService struct {
|
||||
ds dashboards.DashboardService
|
||||
rs rendering.Service
|
||||
}
|
||||
|
||||
func NewBrowserScreenshotService(ds dashboards.DashboardService, rs rendering.Service) ScreenshotService {
|
||||
return &BrowserScreenshotService{
|
||||
ds: ds,
|
||||
rs: rs,
|
||||
}
|
||||
}
|
||||
|
||||
// Take returns a screenshot or an error if either the dashboard does not exist
|
||||
// or it failed to screenshot the dashboard. It uses both the context and the
|
||||
// timeout in ScreenshotOptions, however the timeout in ScreenshotOptions is
|
||||
// sent to the remote browser where it is used as a client timeout.
|
||||
func (s *BrowserScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
q := models.GetDashboardQuery{Uid: opts.DashboardUID}
|
||||
if err := s.ds.GetDashboard(ctx, &q); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
opts = opts.SetDefaults()
|
||||
|
||||
// Compute the URL to screenshot.
|
||||
renderPath := path.Join("d-solo", q.Result.Uid, q.Result.Slug)
|
||||
url := &url.URL{}
|
||||
url.Path = renderPath
|
||||
qParams := url.Query()
|
||||
qParams.Add("orgId", fmt.Sprint(q.Result.OrgId))
|
||||
if opts.PanelID != 0 {
|
||||
qParams.Add("panelId", fmt.Sprint(opts.PanelID))
|
||||
}
|
||||
url.RawQuery = qParams.Encode()
|
||||
path := url.String()
|
||||
|
||||
renderOpts := rendering.Opts{
|
||||
AuthOpts: rendering.AuthOpts{
|
||||
OrgID: q.Result.OrgId,
|
||||
OrgRole: models.ROLE_ADMIN,
|
||||
},
|
||||
ErrorOpts: rendering.ErrorOpts{
|
||||
ErrorConcurrentLimitReached: true,
|
||||
ErrorRenderUnavailable: true,
|
||||
},
|
||||
TimeoutOpts: rendering.TimeoutOpts{
|
||||
Timeout: opts.Timeout,
|
||||
},
|
||||
Width: opts.Width,
|
||||
Height: opts.Height,
|
||||
Theme: opts.Theme,
|
||||
ConcurrentLimit: setting.AlertingRenderLimit,
|
||||
Path: path,
|
||||
}
|
||||
|
||||
result, err := s.rs.Render(ctx, renderOpts, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to take screenshot: %w", err)
|
||||
}
|
||||
|
||||
screenshot := Screenshot{Path: result.FilePath}
|
||||
return &screenshot, nil
|
||||
}
|
||||
|
||||
// CachableScreenshotService caches screenshots.
|
||||
type CachableScreenshotService struct {
|
||||
cache *gocache.Cache
|
||||
service ScreenshotService
|
||||
cacheHits prometheus.Counter
|
||||
cacheMisses prometheus.Counter
|
||||
}
|
||||
|
||||
func NewCachableScreenshotService(r prometheus.Registerer, expiration time.Duration, service ScreenshotService) ScreenshotService {
|
||||
return &CachableScreenshotService{
|
||||
cache: gocache.New(expiration, time.Minute),
|
||||
service: service,
|
||||
cacheHits: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "cache_hits_total",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
cacheMisses: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "cache_misses_total",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Take returns the screenshot from the cache or asks the service to take a
|
||||
// new screenshot and cache it before returning it.
|
||||
func (s *CachableScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
|
||||
|
||||
if v, ok := s.cache.Get(k); ok {
|
||||
defer s.cacheHits.Inc()
|
||||
return v.(*Screenshot), nil
|
||||
}
|
||||
|
||||
defer s.cacheMisses.Inc()
|
||||
screenshot, err := s.service.Take(ctx, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.cache.Set(k, screenshot, 0)
|
||||
|
||||
return screenshot, nil
|
||||
}
|
||||
|
||||
// NoopScreenshotService is a service that takes no-op screenshots.
|
||||
type NoopScreenshotService struct{}
|
||||
|
||||
func (s *NoopScreenshotService) Take(_ context.Context, _ ScreenshotOptions) (*Screenshot, error) {
|
||||
return &Screenshot{}, nil
|
||||
}
|
||||
|
||||
// ObservableScreenshotService is a service that records metrics about screenshots.
|
||||
type ObservableScreenshotService struct {
|
||||
service ScreenshotService
|
||||
duration prometheus.Histogram
|
||||
failures prometheus.Counter
|
||||
successes prometheus.Counter
|
||||
}
|
||||
|
||||
func NewObservableScreenshotService(r prometheus.Registerer, service ScreenshotService) ScreenshotService {
|
||||
return &ObservableScreenshotService{
|
||||
service: service,
|
||||
duration: promauto.With(r).NewHistogram(prometheus.HistogramOpts{
|
||||
Name: "duration_seconds",
|
||||
Buckets: []float64{0.1, 0.25, 0.5, 1, 2, 5, 10, 15},
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
failures: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "failures_total",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
successes: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "successes_total",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *ObservableScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
start := time.Now()
|
||||
defer func() { s.duration.Observe(time.Since(start).Seconds()) }()
|
||||
|
||||
screenshot, err := s.service.Take(ctx, opts)
|
||||
if err != nil {
|
||||
defer s.failures.Inc()
|
||||
} else {
|
||||
defer s.successes.Inc()
|
||||
}
|
||||
return screenshot, err
|
||||
}
|
||||
|
||||
type ScreenshotUnavailableService struct{}
|
||||
|
||||
func (s *ScreenshotUnavailableService) Take(_ context.Context, _ ScreenshotOptions) (*Screenshot, error) {
|
||||
return nil, ErrScreenshotsUnavailable
|
||||
}
|
||||
|
||||
// SingleFlightScreenshotService prevents duplicate screenshots.
|
||||
type SingleFlightScreenshotService struct {
|
||||
f singleflight.Group
|
||||
service ScreenshotService
|
||||
}
|
||||
|
||||
func NewSingleFlightScreenshotService(service ScreenshotService) ScreenshotService {
|
||||
return &SingleFlightScreenshotService{service: service}
|
||||
}
|
||||
|
||||
// Take returns a screenshot or an error. It ensures that at most one screenshot
|
||||
// can be taken at a time for the same dashboard and theme. Duplicate screenshots
|
||||
// wait for the first screenshot to complete and receive the same screenshot.
|
||||
func (s *SingleFlightScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
|
||||
|
||||
v, err, _ := s.f.Do(k, func() (interface{}, error) {
|
||||
return s.service.Take(ctx, opts)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
screenshot := v.(*Screenshot)
|
||||
return screenshot, err
|
||||
}
|
||||
|
||||
// RateLimitScreenshotService ensures that at most N screenshots can be taken
|
||||
// at a time.
|
||||
type RateLimitScreenshotService struct {
|
||||
service ScreenshotService
|
||||
tokens chan struct{}
|
||||
}
|
||||
|
||||
func NewRateLimitScreenshotService(service ScreenshotService, n int64) ScreenshotService {
|
||||
return &RateLimitScreenshotService{
|
||||
service: service,
|
||||
tokens: make(chan struct{}, n),
|
||||
}
|
||||
}
|
||||
|
||||
// Take returns a screenshot or an error. It ensures that at most N screenshots
|
||||
// can be taken at a time. The service has N tokens such that a token is consumed
|
||||
// at the start of a screenshot and returned when the screenshot has either
|
||||
// succeeded or failed. A screenshot can timeout if the context is canceled
|
||||
// while waiting for a token or while the screenshot is being taken.
|
||||
func (s *RateLimitScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
select {
|
||||
// the context is canceled
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
// there is a token available
|
||||
case s.tokens <- struct{}{}:
|
||||
}
|
||||
// acquired token must be returned
|
||||
defer func() {
|
||||
<-s.tokens
|
||||
}()
|
||||
return s.service.Take(ctx, opts)
|
||||
}
|
||||
|
||||
// UploadingScreenshotService uploads taken screenshots.
|
||||
type UploadingScreenshotService struct {
|
||||
service ScreenshotService
|
||||
uploader imguploader.ImageUploader
|
||||
uploadFailures prometheus.Counter
|
||||
uploadSuccesses prometheus.Counter
|
||||
}
|
||||
|
||||
func NewUploadingScreenshotService(r prometheus.Registerer, service ScreenshotService, uploader imguploader.ImageUploader) ScreenshotService {
|
||||
return &UploadingScreenshotService{
|
||||
service: service,
|
||||
uploader: uploader,
|
||||
uploadFailures: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "upload_failures",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
uploadSuccesses: promauto.With(r).NewCounter(prometheus.CounterOpts{
|
||||
Name: "upload_successes",
|
||||
Namespace: namespace,
|
||||
Subsystem: subsystem,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Take uploads a screenshot with a path and returns a new screenshot with the
|
||||
// unmodified path and a URL. It returns the unmodified screenshot on error.
|
||||
func (s *UploadingScreenshotService) Take(ctx context.Context, opts ScreenshotOptions) (*Screenshot, error) {
|
||||
screenshot, err := s.service.Take(ctx, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
url, err := s.uploader.Upload(ctx, screenshot.Path)
|
||||
if err != nil {
|
||||
defer s.uploadFailures.Inc()
|
||||
return screenshot, fmt.Errorf("failed to upload screenshot: %w", err)
|
||||
}
|
||||
screenshot.URL = url
|
||||
|
||||
defer s.uploadSuccesses.Inc()
|
||||
return screenshot, nil
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
package screenshot
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/grafana/grafana/pkg/components/imguploader"
|
||||
"github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/services/dashboards"
|
||||
"github.com/grafana/grafana/pkg/services/rendering"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
)
|
||||
|
||||
func TestScreenshotOptions(t *testing.T) {
|
||||
o := ScreenshotOptions{}
|
||||
assert.Equal(t, ScreenshotOptions{}, o)
|
||||
|
||||
o = o.SetDefaults()
|
||||
assert.Equal(t, ScreenshotOptions{
|
||||
Width: DefaultWidth,
|
||||
Height: DefaultHeight,
|
||||
Theme: DefaultTheme,
|
||||
Timeout: DefaultTimeout,
|
||||
}, o)
|
||||
|
||||
o.Width = 100
|
||||
o = o.SetDefaults()
|
||||
assert.Equal(t, ScreenshotOptions{
|
||||
Width: 100,
|
||||
Height: DefaultHeight,
|
||||
Theme: DefaultTheme,
|
||||
Timeout: DefaultTimeout,
|
||||
}, o)
|
||||
|
||||
o.Height = 100
|
||||
o = o.SetDefaults()
|
||||
assert.Equal(t, ScreenshotOptions{
|
||||
Width: 100,
|
||||
Height: 100,
|
||||
Theme: DefaultTheme,
|
||||
Timeout: DefaultTimeout,
|
||||
}, o)
|
||||
|
||||
o.Theme = "Not a theme"
|
||||
o = o.SetDefaults()
|
||||
assert.Equal(t, ScreenshotOptions{
|
||||
Width: 100,
|
||||
Height: 100,
|
||||
Theme: DefaultTheme,
|
||||
Timeout: DefaultTimeout,
|
||||
}, o)
|
||||
|
||||
o.Timeout = DefaultTimeout + 1
|
||||
assert.Equal(t, ScreenshotOptions{
|
||||
Width: 100,
|
||||
Height: 100,
|
||||
Theme: DefaultTheme,
|
||||
Timeout: DefaultTimeout + 1,
|
||||
}, o)
|
||||
}
|
||||
|
||||
func TestBrowserScreenshotService(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
d := dashboards.FakeDashboardService{}
|
||||
r := rendering.NewMockService(c)
|
||||
s := NewBrowserScreenshotService(&d, r)
|
||||
|
||||
// a non-existent dashboard should return error
|
||||
d.GetDashboardFn = func(ctx context.Context, cmd *models.GetDashboardQuery) error {
|
||||
return models.ErrDashboardNotFound
|
||||
}
|
||||
ctx := context.Background()
|
||||
opts := ScreenshotOptions{}
|
||||
screenshot, err := s.Take(ctx, opts)
|
||||
assert.EqualError(t, err, "Dashboard not found")
|
||||
assert.Nil(t, screenshot)
|
||||
|
||||
d.GetDashboardFn = func(ctx context.Context, cmd *models.GetDashboardQuery) error {
|
||||
cmd.Result = &models.Dashboard{Id: 1, Uid: "foo", Slug: "bar", OrgId: 2}
|
||||
return nil
|
||||
}
|
||||
|
||||
renderOpts := rendering.Opts{
|
||||
AuthOpts: rendering.AuthOpts{
|
||||
OrgID: 2,
|
||||
OrgRole: models.ROLE_ADMIN,
|
||||
},
|
||||
ErrorOpts: rendering.ErrorOpts{
|
||||
ErrorConcurrentLimitReached: true,
|
||||
ErrorRenderUnavailable: true,
|
||||
},
|
||||
TimeoutOpts: rendering.TimeoutOpts{
|
||||
Timeout: DefaultTimeout,
|
||||
},
|
||||
Width: DefaultWidth,
|
||||
Height: DefaultHeight,
|
||||
Theme: DefaultTheme,
|
||||
Path: "d-solo/foo/bar?orgId=2&panelId=4",
|
||||
ConcurrentLimit: setting.AlertingRenderLimit,
|
||||
}
|
||||
|
||||
opts.DashboardUID = "foo"
|
||||
opts.PanelID = 4
|
||||
r.EXPECT().
|
||||
Render(ctx, renderOpts, nil).
|
||||
Return(&rendering.RenderResult{FilePath: "panel.png"}, nil)
|
||||
screenshot, err = s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
|
||||
|
||||
// a timeout should return error
|
||||
r.EXPECT().
|
||||
Render(ctx, renderOpts, nil).
|
||||
Return(nil, rendering.ErrTimeout)
|
||||
screenshot, err = s.Take(ctx, opts)
|
||||
assert.EqualError(t, err, fmt.Sprintf("failed to take screenshot: %s", rendering.ErrTimeout))
|
||||
assert.Nil(t, screenshot)
|
||||
}
|
||||
|
||||
func TestCachableScreenshotService(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
m := NewMockScreenshotService(c)
|
||||
s := NewCachableScreenshotService(prometheus.DefaultRegisterer, time.Second, m)
|
||||
|
||||
ctx := context.Background()
|
||||
opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
|
||||
// should be a miss and ask the mock service to take a screenshot
|
||||
m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "panel.png"}, nil)
|
||||
screenshot, err := s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
|
||||
|
||||
// should be a hit
|
||||
screenshot, err = s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
|
||||
|
||||
// wait 1s and the cached screenshot should have expired
|
||||
<-time.After(time.Second)
|
||||
|
||||
// should be a miss as the cached screenshot has expired
|
||||
m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "panel.png"}, nil)
|
||||
screenshot, err = s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
|
||||
}
|
||||
|
||||
func TestNoopScreenshotService(t *testing.T) {
|
||||
s := NoopScreenshotService{}
|
||||
screenshot, err := s.Take(context.Background(), ScreenshotOptions{})
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, screenshot)
|
||||
}
|
||||
|
||||
func TestScreenshotUnavailableService(t *testing.T) {
|
||||
s := ScreenshotUnavailableService{}
|
||||
screenshot, err := s.Take(context.Background(), ScreenshotOptions{})
|
||||
assert.Equal(t, err, ErrScreenshotsUnavailable)
|
||||
assert.Nil(t, screenshot)
|
||||
}
|
||||
|
||||
func TestSingleFlightScreenshotService(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
m := NewMockScreenshotService(c)
|
||||
s := NewSingleFlightScreenshotService(m)
|
||||
|
||||
ctx := context.Background()
|
||||
opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
|
||||
// expect 1 invocation of the mock service for the same options
|
||||
m.EXPECT().Take(ctx, opts).
|
||||
Do(func(_ context.Context, _ ScreenshotOptions) { <-time.After(time.Second) }).
|
||||
Return(&Screenshot{Path: "panel.png"}, nil)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 5; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel.png"}, *screenshot)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// expect two invocations of the mock service for different dashboards
|
||||
opts1 := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
opts2 := ScreenshotOptions{DashboardUID: "bar", PanelID: 1}
|
||||
m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "foo.png"}, nil)
|
||||
m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "bar.png"}, nil)
|
||||
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "foo.png"}, *screenshot)
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts2)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "bar.png"}, *screenshot)
|
||||
}()
|
||||
wg.Wait()
|
||||
|
||||
// expect two invocations of the mock service for different panels in the same dashboard
|
||||
opts1 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
opts2 = ScreenshotOptions{DashboardUID: "foo", PanelID: 2}
|
||||
m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "panel1.png"}, nil)
|
||||
m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "panel2.png"}, nil)
|
||||
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel1.png"}, *screenshot)
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts2)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "panel2.png"}, *screenshot)
|
||||
}()
|
||||
wg.Wait()
|
||||
|
||||
// expect two invocations of the mock service for different panels in the same dashboard
|
||||
opts1 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1, Theme: models.ThemeDark}
|
||||
opts2 = ScreenshotOptions{DashboardUID: "foo", PanelID: 1, Theme: models.ThemeLight}
|
||||
m.EXPECT().Take(ctx, opts1).Return(&Screenshot{Path: "dark.png"}, nil)
|
||||
m.EXPECT().Take(ctx, opts2).Return(&Screenshot{Path: "light.png"}, nil)
|
||||
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts1)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "dark.png"}, *screenshot)
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
screenshot, err := s.Take(ctx, opts2)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "light.png"}, *screenshot)
|
||||
}()
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestRateLimitScreenshotService(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
m := NewMockScreenshotService(c)
|
||||
s := NewRateLimitScreenshotService(m, 1)
|
||||
ctx := context.Background()
|
||||
opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
|
||||
var v int64
|
||||
for i := 0; i < 10; i++ {
|
||||
m.EXPECT().Take(ctx, opts).
|
||||
Do(func(_ context.Context, _ ScreenshotOptions) {
|
||||
// v should be 0 to show that no tokens have been acquired
|
||||
assert.Equal(t, int64(0), atomic.LoadInt64(&v))
|
||||
atomic.AddInt64(&v, 1)
|
||||
assert.Equal(t, int64(1), atomic.LoadInt64(&v))
|
||||
|
||||
// interrupt so other goroutines can attempt to acquire the token
|
||||
<-time.After(time.Microsecond)
|
||||
|
||||
// v should be 1 to show that no other goroutines acquired the token
|
||||
assert.Equal(t, int64(1), atomic.LoadInt64(&v))
|
||||
atomic.AddInt64(&v, -1)
|
||||
assert.Equal(t, int64(0), atomic.LoadInt64(&v))
|
||||
}).
|
||||
Return(&Screenshot{Path: "foo.png"}, nil)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
result, err := s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{Path: "foo.png"}, *result)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestUploadingScreenshotService(t *testing.T) {
|
||||
c := gomock.NewController(t)
|
||||
defer c.Finish()
|
||||
|
||||
m := NewMockScreenshotService(c)
|
||||
u := imguploader.NewMockImageUploader(c)
|
||||
s := NewUploadingScreenshotService(prometheus.DefaultRegisterer, m, u)
|
||||
|
||||
ctx := context.Background()
|
||||
opts := ScreenshotOptions{DashboardUID: "foo", PanelID: 1}
|
||||
|
||||
m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "foo.png"}, nil)
|
||||
u.EXPECT().Upload(ctx, "foo.png").Return("https://example.com/foo.png", nil)
|
||||
screenshot, err := s.Take(ctx, opts)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, Screenshot{
|
||||
Path: "foo.png",
|
||||
URL: "https://example.com/foo.png",
|
||||
}, *screenshot)
|
||||
|
||||
// error on upload should still return screenshot on disk
|
||||
m.EXPECT().Take(ctx, opts).Return(&Screenshot{Path: "foo.png"}, nil)
|
||||
u.EXPECT().Upload(ctx, "foo.png").Return("", errors.New("service is unavailable"))
|
||||
screenshot, err = s.Take(ctx, opts)
|
||||
assert.EqualError(t, err, "failed to upload screenshot: service is unavailable")
|
||||
assert.Equal(t, Screenshot{
|
||||
Path: "foo.png",
|
||||
}, *screenshot)
|
||||
}
|
||||
Reference in New Issue
Block a user