internal/screentest: add support for writing test output to GCS

Test scripts parse output locations that begin with gs:// as Cloud Storage buckets. When a bucket is used diff images and cached screenshots are written to and read from GCS objects. Change-Id: I985ccf301ada1bfde82e4e61e1ddf724a824fcb6 Reviewed-on: https://go-review.googlesource.com/c/website/+/373720 Run-TryBot: Jamal Carvalho <jamal@golang.org> Trust: Jamal Carvalho <jamalcarvalho@google.com> Reviewed-by: Jonathan Amsterdam <jba@google.com> TryBot-Result: Gopher Robot <gobot@golang.org>
2021-12-24 00:36:24 +00:00 · 2021-12-24 00:36:24 +00:00 · 995577f14f
--- a/go.mod
+++ b/go.mod
@ -5,6 +5,7 @@ go 1.16
 require (
 	cloud.google.com/go v0.88.0
 	cloud.google.com/go/datastore v1.2.0
+	cloud.google.com/go/storage v1.10.0
 	github.com/chromedp/cdproto v0.0.0-20211205231339-d2673e93eee4
 	github.com/chromedp/chromedp v0.7.6
 	github.com/evanw/esbuild v0.14.7
--- a/go.sum
+++ b/go.sum
@ -44,6 +44,7 @@ cloud.google.com/go/storage v1.0.0/go.mod h1:IhtSnM/ZTZV8YYJWCY8RULGVqBDmpoyjwiy
 cloud.google.com/go/storage v1.5.0/go.mod h1:tpKbwo567HUNpVclU5sGELwQWBDZ8gh0ZeosJ0Rtdos=
 cloud.google.com/go/storage v1.6.0/go.mod h1:N7U0C8pVQ/+NIKOBQyamJIeKQKkZ+mxpohlUTyfDhBk=
 cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RXyy7KQOVs=
+cloud.google.com/go/storage v1.10.0 h1:STgFzyU5/8miMl0//zKh2aQeTyeaUH3WN9bSUiJ09bA=
 cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0=
 contrib.go.opencensus.io/exporter/prometheus v0.3.0/go.mod h1:rpCPVQKhiyH8oomWgm34ZmgIdZa8OVYO5WAIygPbBBE=
 contrib.go.opencensus.io/exporter/stackdriver v0.13.5/go.mod h1:aXENhDJ1Y4lIg4EUaVTwzvYETVNZk10Pu26tevFKLUc=
@ -324,9 +325,11 @@ github.com/google/go-github v17.0.0+incompatible/go.mod h1:zLgOLi98H3fifZn+44m+u
 github.com/google/go-github/v35 v35.2.0/go.mod h1:s0515YVTI+IMrDoy9Y4pHt9ShGpzHvHO8rZ7L7acgvs=
 github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO6wN/zVPAxq5ck=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/martian v2.1.0+incompatible h1:/CP5g8u/VJHijgedC/Legn3BAbAaWPgecwXBIDzw5no=
 github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXiIaQGbYVAs8BPL6v8lEs=
 github.com/google/martian/v3 v3.0.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=
 github.com/google/martian/v3 v3.1.0/go.mod h1:y5Zk1BBys9G+gd6Jrk0W3cC1+ELVxBWuIGO+w/tUAp0=
+github.com/google/martian/v3 v3.2.1 h1:d8MncMlErDFTwQGBK1xhv026j9kqhvw1Qv9IbWT1VLQ=
 github.com/google/martian/v3 v3.2.1/go.mod h1:oBOf6HBosgwRXnUGWUB05QECsc6uvmMiJ3+6W4l/CUk=
 github.com/google/pprof v0.0.0-20181206194817-3ea8567a2e57/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
 github.com/google/pprof v0.0.0-20190515194954-54271f7e092f/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc=
@ -580,6 +583,7 @@ github.com/openzipkin-contrib/zipkin-go-opentracing v0.4.5/go.mod h1:/wsWhb9smxS
 github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJc5AZX7/PBEpw=
 github.com/openzipkin/zipkin-go v0.2.1/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4=
 github.com/openzipkin/zipkin-go v0.2.2/go.mod h1:NaW6tEwdmWMaCDZzg8sh+IBNOxHMPnhQw8ySjnjRyN4=
+github.com/orisano/pixelmatch v0.0.0-20210112091706-4fa4c7ba91d5 h1:1SoBaSPudixRecmlHXb/GxmaD3fLMtHIDN13QujwQuc=
 github.com/orisano/pixelmatch v0.0.0-20210112091706-4fa4c7ba91d5/go.mod h1:nZgzbfBr3hhjoZnS66nKrHmduYNpc34ny7RK4z5/HM0=
 github.com/pact-foundation/pact-go v1.0.4/go.mod h1:uExwJY4kCzNPcHRj+hCR/HBbOOIwwtUjcrb0b5/5kLM=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
@ -936,7 +940,6 @@ golang.org/x/sys v0.0.0-20210514084401-e8d321eab015/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210603125802-9665404d3644/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c h1:F1jZWGFhYfh0Ci55sIpILtKKK8p3i2/krTr0H1rg74I=
 golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210908233432-aa78b53d3365/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211124211545-fe61309f8881 h1:TyHqChC80pFkXWraUUf6RuB5IqFdQieMLwwCJokV2pc=
--- a/internal/screentest/screentest.go
+++ b/internal/screentest/screentest.go
@ -46,6 +46,10 @@
 //
 //  output testdata/snapshots
 //
+// USE output BUCKETNAME for screentest to upload test output to a Cloud Storage bucket.
+//
+//  output gs://bucket-name
+//
 // Use test NAME to create a name for the testcase.
 //
 //  test about page
@ -100,7 +104,9 @@ import (
 	"fmt"
 	"image"
 	"image/png"
+	"io"
 	"io/fs"
+	"io/ioutil"
 	"log"
 	"net/url"
 	"os"
@ -112,6 +118,7 @@ import (
 	"text/template"
 	"time"

+	"cloud.google.com/go/storage"
 	"github.com/chromedp/cdproto/network"
 	"github.com/chromedp/cdproto/page"
 	"github.com/chromedp/chromedp"
@ -203,6 +210,7 @@ const (
 	browserWidth  = 1536
 	browserHeight = 960
 	cacheSuffix   = "::cache"
+	gcsScheme     = "gs://"
 )

 var sanitize = regexp.MustCompile("[.*<>?`'|/\\: ]")
@ -221,6 +229,7 @@ type testcase struct {
 	urlA, urlB                string
 	headers                   map[string]interface{}
 	cacheA, cacheB            bool
+	gcsBucket                 bool
 	outImgA, outImgB, outDiff string
 	viewportWidth             int
 	viewportHeight            int
@ -249,6 +258,7 @@ func readTests(file string, vars map[string]string) ([]*testcase, error) {
 		originA, originB   string
 		headers            map[string]interface{}
 		cacheA, cacheB     bool
+		gcsBucket          bool
 		width, height      int
 		lineNo             int
 	)
@ -308,6 +318,9 @@ func readTests(file string, vars map[string]string) ([]*testcase, error) {
 			}
 			headers[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
 		case "OUTPUT":
+			if strings.HasPrefix(args, gcsScheme) {
+				gcsBucket = true
+			}
 			out, err = outDir(args, "")
 			if err != nil {
 				return nil, fmt.Errorf("outDir(%q, %q): %w", args, file, err)
@ -375,6 +388,7 @@ func readTests(file string, vars map[string]string) ([]*testcase, error) {
 				viewportHeight: height,
 				cacheA:         cacheA,
 				cacheB:         cacheB,
+				gcsBucket:      gcsBucket,
 			}
 			tests = append(tests, test)
 			field, args := splitOneField(args)
@ -399,6 +413,9 @@ func readTests(file string, vars map[string]string) ([]*testcase, error) {
 				test.screenshotElement = args
 			}
 			outfile := filepath.Join(out, sanitized(test.name))
+			if gcsBucket {
+				outfile = out + "/" + sanitized(test.name)
+			}
 			test.outImgA = outfile + "." + sanitized(urlA.Host) + ".png"
 			test.outImgB = outfile + "." + sanitized(urlB.Host) + ".png"
 			test.outDiff = outfile + ".diff.png"
@ -415,6 +432,9 @@ func readTests(file string, vars map[string]string) ([]*testcase, error) {

 // outDir prepares a diff output directory for a given testfile.
 func outDir(dir, testfile string) (string, error) {
+	if strings.HasPrefix(dir, gcsScheme) {
+		return dir, nil
+	}
 	if testfile != "" {
 		dir = filepath.Join(dir, sanitized(filepath.Base(testfile)))
 	}
@ -488,17 +508,17 @@ func runDiff(ctx context.Context, test *testcase, update bool) (err error) {
 	fmt.Printf("%s != %s (%s)\n", test.urlA, test.urlB, since)
 	g = &errgroup.Group{}
 	g.Go(func() error {
-		return writePNG(&result.Image, test.outDiff)
+		return writePNG(test, &result.Image, test.outDiff)
 	})
 	// Only write screenshots if they haven't already been written to the cache.
 	if !test.cacheA {
 		g.Go(func() error {
-			return writePNG(screenA, test.outImgA)
+			return writePNG(test, screenA, test.outImgA)
 		})
 	}
 	if !test.cacheB {
 		g.Go(func() error {
-			return writePNG(screenB, test.outImgB)
+			return writePNG(test, screenB, test.outImgB)
 		})
 	}
 	if err := g.Wait(); err != nil {
@ -516,35 +536,48 @@ func screenshot(ctx context.Context, test *testcase,
 ) (_ *image.Image, err error) {
 	var data []byte
 	// If cache is enabled, try to read the file from the cache.
-	if cache {
+	if cache && test.gcsBucket {
+		client, err := storage.NewClient(ctx)
+		if err != nil {
+			return nil, fmt.Errorf("storage.NewClient(err): %w", err)
+		}
+		bkt, obj := gcsParts(file)
+		r, err := client.Bucket(bkt).Object(obj).NewReader(ctx)
+		if err != nil && !errors.Is(err, storage.ErrObjectNotExist) {
+			return nil, fmt.Errorf("object.NewReader(ctx): %w", err)
+		} else if err == nil {
+			defer r.Close()
+			data, err = ioutil.ReadAll(r)
+			if err != nil {
+				return nil, fmt.Errorf("ioutil.ReadAll(...): %w", err)
+			}
+		}
+	} else if cache {
 		data, err = os.ReadFile(file)
-		if errors.Is(err, fs.ErrNotExist) {
-			// If screenshot is not found, this must be the first time the test has run.
-			// Set update to true to capture a new screenshot.
-			update = true
-		} else if err != nil {
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
 			return nil, fmt.Errorf("os.ReadFile(...): %w", err)
 		}
 	}
-	// If cache is false, this is the first test run, or an update is requested
+	// If cache is false, an update is requested, or this is the first test run
 	// we capture a new screenshot from a live URL.
-	if !cache || update {
+	if !cache || update || data == nil {
+		update = true
 		data, err = captureScreenshot(ctx, test, url)
 		if err != nil {
 			return nil, fmt.Errorf("captureScreenshot(ctx, %q, %q): %w", url, test, err)
 		}
 	}
-	// Write to the cache.
-	if cache && update {
-		if err := os.WriteFile(file, data, os.ModePerm); err != nil {
-			return nil, fmt.Errorf("os.WriteFile(...): %w", err)
-		}
-		fmt.Printf("updated %s\n", file)
-	}
 	img, _, err := image.Decode(bytes.NewReader(data))
 	if err != nil {
 		return nil, fmt.Errorf("image.Decode(...): %w", err)
 	}
+	// Write to the cache.
+	if cache && update {
+		if err := writePNG(test, &img, file); err != nil {
+			return nil, fmt.Errorf("os.WriteFile(...): %w", err)
+		}
+		fmt.Printf("updated %s\n", file)
+	}
 	return &img, nil
 }

@ -581,10 +614,21 @@ func captureScreenshot(ctx context.Context, test *testcase, url string) ([]byte,
 }

 // writePNG writes image data to a png file.
-func writePNG(i *image.Image, filename string) error {
-	f, err := os.Create(filename)
-	if err != nil {
-		return fmt.Errorf("os.Create(%q): %w", filename, err)
+func writePNG(test *testcase, i *image.Image, filename string) (err error) {
+	var f io.WriteCloser
+	if test.gcsBucket {
+		ctx := context.Background()
+		client, err := storage.NewClient(ctx)
+		if err != nil {
+			return fmt.Errorf("storage.NewClient(ctx): %w", err)
+		}
+		bkt, obj := gcsParts(filename)
+		f = client.Bucket(bkt).Object(obj).NewWriter(ctx)
+	} else {
+		f, err = os.Create(filename)
+		if err != nil {
+			return fmt.Errorf("os.Create(%q): %w", filename, err)
+		}
 	}
 	err = png.Encode(f, *i)
 	if err != nil {
@ -604,6 +648,16 @@ func sanitized(text string) string {
 	return sanitize.ReplaceAllString(text, "-")
 }

+// gcsParts splits a Cloud Storage filename into bucket name and
+// object name parts.
+func gcsParts(filename string) (bucket, object string) {
+	filename = strings.TrimPrefix(filename, gcsScheme)
+	n := strings.Index(filename, "/")
+	bucket = filename[:n]
+	object = filename[n+1:]
+	return bucket, object
+}
+
 // waitForEvent waits for browser lifecycle events. This is useful for
 // ensuring the page is fully loaded before capturing screenshots.
 func waitForEvent(eventName string) chromedp.ActionFunc {
--- a/internal/screentest/screentest_test.go
+++ b/internal/screentest/screentest_test.go
@ -143,6 +143,20 @@ func TestReadTests(t *testing.T) {
 						chromedp.Evaluate("console.log('Hello, world!')", nil),
 					},
 				},
+				{
+					name:           "gcs-output",
+					urlA:           "https://pkg.go.dev/gcs-output",
+					cacheA:         true,
+					urlB:           "http://localhost:8080/gcs-output",
+					gcsBucket:      true,
+					headers:        map[string]interface{}{"Authorization": "Bearer token"},
+					outImgA:        "gs://bucket-name/gcs-output.pkg-go-dev.png",
+					outImgB:        "gs://bucket-name/gcs-output.localhost-8080.png",
+					outDiff:        "gs://bucket-name/gcs-output.diff.png",
+					screenshotType: viewportScreenshot,
+					viewportWidth:  1536,
+					viewportHeight: 960,
+				},
 			},
 			wantErr: false,
 		},
@ -286,3 +300,41 @@ func headerServer() error {
 	})
 	return http.ListenAndServe(fmt.Sprintf(":%d", 6061), mux)
 }
+
+func Test_gcsParts(t *testing.T) {
+	type args struct {
+		filename string
+	}
+	tests := []struct {
+		name       string
+		args       args
+		wantBucket string
+		wantObject string
+	}{
+		{
+			args: args{
+				filename: "gs://bucket-name/object-name",
+			},
+			wantBucket: "bucket-name",
+			wantObject: "object-name",
+		},
+		{
+			args: args{
+				filename: "gs://bucket-name/subdir/object-name",
+			},
+			wantBucket: "bucket-name",
+			wantObject: "subdir/object-name",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotBucket, gotObject := gcsParts(tt.args.filename)
+			if gotBucket != tt.wantBucket {
+				t.Errorf("gcsParts() gotBucket = %v, want %v", gotBucket, tt.wantBucket)
+			}
+			if gotObject != tt.wantObject {
+				t.Errorf("gcsParts() gotObject = %v, want %v", gotObject, tt.wantObject)
+			}
+		})
+	}
+}
--- a/internal/screentest/testdata/readtests.txt
+++ b/internal/screentest/testdata/readtests.txt
@ -33,3 +33,7 @@ test eval
 pathname /eval
 eval console.log('hello, world!')
 capture
+
+output gs://bucket-name
+pathname /gcs-output
+capture