Wire up the local and remote persisters

browser/file_persister.go

@@ -0,0 +1,100 @@
+package browser
+
+import (
+	"errors"
+	"fmt"
+	"net/url"
+	"strings"
+
+	"github.com/grafana/xk6-browser/env"
+	"github.com/grafana/xk6-browser/storage"
+)
+
+type presignedURLConfig struct {
+	getterURL string
+	headers   map[string]string
+	basePath  string
+}
+
+// newScreenshotPersister will return either a persister that persists file to the local
+// disk or uploads the files to a remote location. This decision depends on whether
+// the K6_BROWSER_SCREENSHOTS_OUTPUT env var is setup with the correct configs.
+func newScreenshotPersister(envLookup env.LookupFunc) (filePersister, error) {
+	envVar, ok := envLookup(env.ScreenshotsOutput)
+	if !ok || envVar == "" {
+		return &storage.LocalFilePersister{}, nil
+	}
+
+	popts, err := parsePresignedURLEnvVar(envVar)
+	if err != nil {
+		return nil, fmt.Errorf("parsing %s: %w", env.ScreenshotsOutput, err)
+	}
+
+	return storage.NewRemoteFilePersister(popts.getterURL, popts.headers, popts.basePath), nil
+}
+
+// parsePresignedURLEnvVar will parse a value such as:
+// url=https://127.0.0.1/,basePath=/screenshots,header.1=a,header.2=b
+// and return them.
+//
+//nolint:cyclop
+func parsePresignedURLEnvVar(envVarValue string) (presignedURLConfig, error) {
+	ss := strings.Split(envVarValue, ",")
+
+	presignedURL := presignedURLConfig{
+		headers: make(map[string]string),
+	}
+	for _, s := range ss {
+		// The key value pair should be of the form key=value, so split
+		// on '=' to retrieve the key and value separately.
+		kv := strings.Split(s, "=")
+		if len(kv) != 2 {
+			return presignedURLConfig{}, fmt.Errorf("format of value must be k=v, received %q", s)
+		}
+
+		k := kv[0]
+		v := kv[1]
+
+		// A key with "header." means that the header name is encoded in the
+		// key, separated by a ".". Split the header on "." to retrieve the
+		// header name. The header value should be present in v from the previous
+		// split.
+		var hv, hk string
+		if strings.Contains(k, "header.") {
+			hv = v
+
+			hh := strings.Split(k, ".")
+			if len(hh) != 2 {
+				return presignedURLConfig{}, fmt.Errorf("format of header must be header.k=v, received %q", s)
+			}
+
+			k = hh[0]
+			hk = hh[1]
+
+			if hk == "" {
+				return presignedURLConfig{}, fmt.Errorf("empty header key, received %q", s)
+			}
+		}
+
+		switch k {
+		case "url":
+			u, err := url.ParseRequestURI(v)
+			if err != nil && u.Scheme != "" && u.Host != "" {
+				return presignedURLConfig{}, fmt.Errorf("invalid url %q", s)
+			}
+			presignedURL.getterURL = v
+		case "basePath":
+			presignedURL.basePath = v
+		case "header":
+			presignedURL.headers[hk] = hv
+		default:
+			return presignedURLConfig{}, fmt.Errorf("invalid option %q", k)
+		}
+	}
+
+	if presignedURL.getterURL == "" {
+		return presignedURLConfig{}, errors.New("missing required url")
+	}
+
+	return presignedURL, nil
+}

browser/file_persister_test.go

@@ -0,0 +1,174 @@
+package browser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/grafana/xk6-browser/env"
+	"github.com/grafana/xk6-browser/storage"
+)
+
+func Test_newScreenshotPersister(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		envLookup env.LookupFunc
+		wantType  filePersister
+		wantErr   bool
+	}{
+		{
+			name:      "local_no_env_var",
+			envLookup: env.EmptyLookup,
+			wantType:  &storage.LocalFilePersister{},
+		},
+		{
+			name: "local_empty_env_var",
+			envLookup: env.ConstLookup(
+				env.ScreenshotsOutput,
+				"",
+			),
+			wantType: &storage.LocalFilePersister{},
+		},
+		{
+			name: "remote",
+			envLookup: env.ConstLookup(
+				env.ScreenshotsOutput,
+				"url=https://127.0.0.1/,basePath=/screenshots,header.1=a",
+			),
+			wantType: &storage.RemoteFilePersister{},
+		},
+		{
+			name: "remote_parse_failed",
+			envLookup: env.ConstLookup(
+				env.ScreenshotsOutput,
+				"basePath=/screenshots,header.1=a",
+			),
+			wantErr: true,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			gotType, err := newScreenshotPersister(tt.envLookup)
+			if tt.wantErr {
+				assert.Error(t, err)
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.IsType(t, tt.wantType, gotType)
+		})
+	}
+}
+
+func Test_parsePresignedURLEnvVar(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name        string
+		envVarValue string
+		want        presignedURLConfig
+		wantErr     string
+	}{
+		{
+			name:        "url_headers_basePath",
+			envVarValue: "url=https://127.0.0.1/,basePath=/screenshots,header.1=a,header.2=b",
+			want: presignedURLConfig{
+				getterURL: "https://127.0.0.1/",
+				basePath:  "/screenshots",
+				headers: map[string]string{
+					"1": "a",
+					"2": "b",
+				},
+			},
+		},
+		{
+			name:        "url_headers",
+			envVarValue: "url=https://127.0.0.1/,header.1=a,header.2=b",
+			want: presignedURLConfig{
+				getterURL: "https://127.0.0.1/",
+				headers: map[string]string{
+					"1": "a",
+					"2": "b",
+				},
+			},
+		},
+		{
+			name:        "url",
+			envVarValue: "url=https://127.0.0.1/",
+			want: presignedURLConfig{
+				getterURL: "https://127.0.0.1/",
+				headers:   map[string]string{},
+			},
+		},
+		{
+			name:        "url_basePath",
+			envVarValue: "url=https://127.0.0.1/,basePath=/screenshots",
+			want: presignedURLConfig{
+				getterURL: "https://127.0.0.1/",
+				basePath:  "/screenshots",
+				headers:   map[string]string{},
+			},
+		},
+		{
+			name:        "empty_basePath",
+			envVarValue: "url=https://127.0.0.1/,basePath=",
+			want: presignedURLConfig{
+				getterURL: "https://127.0.0.1/",
+				basePath:  "",
+				headers:   map[string]string{},
+			},
+		},
+		{
+			name:        "empty",
+			envVarValue: "",
+			wantErr:     `format of value must be k=v, received ""`,
+		},
+		{
+			name:        "missing_url",
+			envVarValue: "basePath=/screenshots,header.1=a,header.2=b",
+			wantErr:     "missing required url",
+		},
+		{
+			name:        "invalid_option",
+			envVarValue: "ulr=https://127.0.0.1/",
+			wantErr:     "invalid option",
+		},
+		{
+			name:        "empty_header_key",
+			envVarValue: "url=https://127.0.0.1/,header.=a",
+			wantErr:     "empty header key",
+		},
+		{
+			name:        "invalid_format",
+			envVarValue: "url==https://127.0.0.1/",
+			wantErr:     "format of value must be k=v",
+		},
+		{
+			name:        "invalid_header_format",
+			envVarValue: "url=https://127.0.0.1/,header..asd=a",
+			wantErr:     "format of header must be header.k=v",
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := parsePresignedURLEnvVar(tt.envVarValue)
+			if tt.wantErr != "" {
+				assert.ErrorContains(t, err, tt.wantErr)
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

browser/mapping.go

@@ -284,7 +284,7 @@ func mapElementHandle(vu moduleVU, eh *common.ElementHandle) mapping {
 				return nil, fmt.Errorf("parsing frame screenshot options: %w", err)
 			}
 
-			bb, err := eh.Screenshot(popts, vu.LocalFilePersister)
+			bb, err := eh.Screenshot(popts, vu.filePersister)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
 			}
@@ -686,7 +686,7 @@ func mapPage(vu moduleVU, p *common.Page) mapping {
 				return nil, fmt.Errorf("parsing page screenshot options: %w", err)
 			}
 
-			bb, err := p.Screenshot(popts, vu.LocalFilePersister)
+			bb, err := p.Screenshot(popts, vu.filePersister)
 			if err != nil {
 				return nil, err //nolint:wrapcheck
 			}

browser/module.go

@@ -8,6 +8,7 @@ package browser
 
 import (
 	"context"
+	"io"
 	"log"
 	"net/http"
 	_ "net/http/pprof" //nolint:gosec
@@ -18,19 +19,25 @@ import (
 	"github.com/grafana/xk6-browser/common"
 	"github.com/grafana/xk6-browser/env"
 	"github.com/grafana/xk6-browser/k6ext"
-	"github.com/grafana/xk6-browser/storage"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
 
 type (
+	// filePersister is the type that all file persisters must implement. It's job is
+	// to persist a file somewhere, hiding the details of where and how from the caller.
+	filePersister interface {
+		Persist(ctx context.Context, path string, data io.Reader) (err error)
+	}
+
 	// RootModule is the global module instance that will create module
 	// instances for each VU.
 	RootModule struct {
 		PidRegistry    *pidRegistry
 		remoteRegistry *remoteRegistry
 		initOnce       *sync.Once
 		tracesMetadata map[string]string
+		filePersister  filePersister
 	}
 
 	// JSModule exposes the properties available to the JS script.
@@ -82,8 +89,8 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 					m.PidRegistry,
 					m.tracesMetadata,
 				),
-				taskQueueRegistry:  newTaskQueueRegistry(vu),
-				LocalFilePersister: &storage.LocalFilePersister{},
+				taskQueueRegistry: newTaskQueueRegistry(vu),
+				filePersister:     m.filePersister,
 			}),
 			Devices:         common.GetDevices(),
 			NetworkProfiles: common.GetNetworkProfiles(),
@@ -115,6 +122,10 @@ func (m *RootModule) initialize(vu k6modules.VU) {
 	if _, ok := initEnv.LookupEnv(env.EnableProfiling); ok {
 		go startDebugServer()
 	}
+	m.filePersister, err = newScreenshotPersister(initEnv.LookupEnv)
+	if err != nil {
+		k6ext.Abort(vu.Context(), "failed to create file persister: %v", err)
+	}
 }
 
 func startDebugServer() {

browser/modulevu.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/grafana/xk6-browser/common"
 	"github.com/grafana/xk6-browser/k6ext"
-	"github.com/grafana/xk6-browser/storage"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
@@ -22,7 +21,7 @@ type moduleVU struct {
 
 	*taskQueueRegistry
 
-	*storage.LocalFilePersister
+	filePersister
 }
 
 // browser returns the VU browser instance for the current iteration.

common/element_handle.go

@@ -17,7 +17,6 @@ import (
 
 	"github.com/grafana/xk6-browser/common/js"
 	"github.com/grafana/xk6-browser/k6ext"
-	"github.com/grafana/xk6-browser/storage"
 )
 
 const resultDone = "done"
@@ -1180,7 +1179,7 @@ func (h *ElementHandle) setChecked(apiCtx context.Context, checked bool, p *Posi
 // Screenshot will instruct Chrome to save a screenshot of the current element and save it to specified file.
 func (h *ElementHandle) Screenshot(
 	opts *ElementHandleScreenshotOptions,
-	fp *storage.LocalFilePersister,
+	sp ScreenshotPersister,
 ) ([]byte, error) {
 	spanCtx, span := TraceAPICall(
 		h.ctx,
@@ -1191,7 +1190,7 @@ func (h *ElementHandle) Screenshot(
 
 	span.SetAttributes(attribute.String("screenshot.path", opts.Path))
 
-	s := newScreenshotter(spanCtx, fp)
+	s := newScreenshotter(spanCtx, sp)
 	buf, err := s.screenshotElement(h, opts)
 	if err != nil {
 		return nil, fmt.Errorf("taking screenshot of elementHandle: %w", err)