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>
diff --git a/go.mod b/go.mod
index 214c029..136b1df 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@
 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
diff --git a/go.sum b/go.sum
index 31162e4..875940a 100644
--- a/go.sum
+++ b/go.sum
@@ -44,6 +44,7 @@
 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/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/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-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=
diff --git a/internal/screentest/screentest.go b/internal/screentest/screentest.go
index 17ee7cf..e6965d4 100644
--- 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 @@
 	"fmt"
 	"image"
 	"image/png"
+	"io"
 	"io/fs"
+	"io/ioutil"
 	"log"
 	"net/url"
 	"os"
@@ -112,6 +118,7 @@
 	"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 @@
 	browserWidth  = 1536
 	browserHeight = 960
 	cacheSuffix   = "::cache"
+	gcsScheme     = "gs://"
 )
 
 var sanitize = regexp.MustCompile("[.*<>?`'|/\\: ]")
@@ -221,6 +229,7 @@
 	urlA, urlB                string
 	headers                   map[string]interface{}
 	cacheA, cacheB            bool
+	gcsBucket                 bool
 	outImgA, outImgB, outDiff string
 	viewportWidth             int
 	viewportHeight            int
@@ -249,6 +258,7 @@
 		originA, originB   string
 		headers            map[string]interface{}
 		cacheA, cacheB     bool
+		gcsBucket          bool
 		width, height      int
 		lineNo             int
 	)
@@ -308,6 +318,9 @@
 			}
 			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 @@
 				viewportHeight: height,
 				cacheA:         cacheA,
 				cacheB:         cacheB,
+				gcsBucket:      gcsBucket,
 			}
 			tests = append(tests, test)
 			field, args := splitOneField(args)
@@ -399,6 +413,9 @@
 				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 @@
 
 // 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 @@
 	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 @@
 ) (_ *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 @@
 }
 
 // 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 @@
 	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 {
diff --git a/internal/screentest/screentest_test.go b/internal/screentest/screentest_test.go
index 8d68f0c..1cc98f0 100644
--- a/internal/screentest/screentest_test.go
+++ b/internal/screentest/screentest_test.go
@@ -143,6 +143,20 @@
 						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 @@
 	})
 	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)
+			}
+		})
+	}
+}
diff --git a/internal/screentest/testdata/readtests.txt b/internal/screentest/testdata/readtests.txt
index d1f1f7a..8f1ec96 100644
--- a/internal/screentest/testdata/readtests.txt
+++ b/internal/screentest/testdata/readtests.txt
@@ -33,3 +33,7 @@
 pathname /eval
 eval console.log('hello, world!')
 capture
+
+output gs://bucket-name
+pathname /gcs-output
+capture