{cmd,internal}/screentest: add support for templates

Screentest scripts are parsed as go templates. Passing
headers from the command line is replaced by the header
keyword in test scripts. Data can be interpolated into
headers using go template strings.

For example `header Authorization: Bearer {{.Token}}`
would become `header Authorization: Bearer abcdef` after
running `screentest -vars "Token:abcdef" testdata/file.txt`.

Change-Id: Ia27b9e12c68ca3bdbe609d3714d57d8867c2b351
Reviewed-on: https://go-review.googlesource.com/c/website/+/373715
Trust: Jamal Carvalho <jamalcarvalho@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
diff --git a/cmd/screentest/main.go b/cmd/screentest/main.go
index 0174aa4..0476314 100644
--- a/cmd/screentest/main.go
+++ b/cmd/screentest/main.go
@@ -2,8 +2,7 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Command screentest runs the visual diff check for the set of scripts
-// provided by the flag -testdata.
+// Command screentest runs the screentest check for a set of scripts.
 package main
 
 import (
@@ -17,8 +16,8 @@
 )
 
 var (
-	update  = flag.Bool("update", false, "update cached screenshots")
-	headers = flag.String("H", "", "set request headers")
+	update = flag.Bool("update", false, "update cached screenshots")
+	vars   = flag.String("vars", "", "provide variables to the script template as comma separated KEY:VALUE pairs")
 )
 
 func main() {
@@ -32,17 +31,17 @@
 		flag.Usage()
 		os.Exit(1)
 	}
-	hdr := make(map[string]interface{})
-	if *headers != "" {
-		for _, h := range strings.Split(*headers, ",") {
-			parts := strings.Split(h, ":")
+	parsedVars := make(map[string]string)
+	if *vars != "" {
+		for _, pair := range strings.Split(*vars, ",") {
+			parts := strings.SplitN(pair, ":", 2)
 			if len(parts) != 2 {
-				log.Fatalf("invalid header %s", h)
+				log.Fatal(fmt.Errorf("invalid key value pair, %q", pair))
 			}
-			hdr[parts[0]] = parts[1]
+			parsedVars[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
 		}
 	}
-	if err := screentest.CheckHandler(args[0], *update, hdr); err != nil {
+	if err := screentest.CheckHandler(args[0], *update, parsedVars); err != nil {
 		log.Fatal(err)
 	}
 }
diff --git a/internal/screentest/screentest.go b/internal/screentest/screentest.go
index 4b97b20..e484488 100644
--- a/internal/screentest/screentest.go
+++ b/internal/screentest/screentest.go
@@ -7,13 +7,13 @@
 //
 // Scripts
 //
-// A script is a text file containing a sequence of testcases, separated by
+// A script is a template file containing a sequence of testcases, separated by
 // blank lines. Lines beginning with # characters are ignored as comments. A
 // testcase is a sequence of lines describing actions to take on a page, along
 // with the dimensions of the screenshots to be compared. For example, here is
 // a trivial script:
 //
-//  compare https://go.dev http://localhost:6060
+//  compare https://go.dev {{.ComparisonURL}}
 //  pathname /about
 //  capture fullscreen
 //
@@ -33,6 +33,10 @@
 //
 //  compare https://go.dev http://localhost:6060
 //
+// Use header KEY:VALUE to add headers to requests
+//
+//  header Authorization: Bearer token
+//
 // Add the ::cache suffix to cache the images from an origin for subsequent
 // test runs.
 //
@@ -105,6 +109,7 @@
 	"strconv"
 	"strings"
 	"testing"
+	"text/template"
 	"time"
 
 	"github.com/chromedp/cdproto/network"
@@ -116,7 +121,7 @@
 
 // CheckHandler runs the test scripts matched by glob. If any errors are
 // encountered, CheckHandler returns an error listing the problems.
-func CheckHandler(glob string, update bool, headers map[string]interface{}) error {
+func CheckHandler(glob string, update bool, vars map[string]string) error {
 	ctx := context.Background()
 	files, err := filepath.Glob(glob)
 	if err != nil {
@@ -132,7 +137,7 @@
 	defer cancel()
 	var buf bytes.Buffer
 	for _, file := range files {
-		tests, err := readTests(file)
+		tests, err := readTests(file, vars)
 		if err != nil {
 			return fmt.Errorf("readTestdata(%q): %w", file, err)
 		}
@@ -143,7 +148,7 @@
 		defer cancel()
 		var hdr bool
 		for _, test := range tests {
-			if err := runDiff(ctx, test, update, headers); err != nil {
+			if err := runDiff(ctx, test, update); err != nil {
 				if !hdr {
 					fmt.Fprintf(&buf, "%s\n\n", file)
 					hdr = true
@@ -160,7 +165,7 @@
 }
 
 // TestHandler runs the test script files matched by glob.
-func TestHandler(t *testing.T, glob string, update bool, headers map[string]interface{}) error {
+func TestHandler(t *testing.T, glob string, update bool, vars map[string]string) error {
 	ctx := context.Background()
 	files, err := filepath.Glob(glob)
 	if err != nil {
@@ -175,7 +180,7 @@
 	)...)
 	defer cancel()
 	for _, file := range files {
-		tests, err := readTests(file)
+		tests, err := readTests(file, vars)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -183,7 +188,7 @@
 		defer cancel()
 		for _, test := range tests {
 			t.Run(test.name, func(t *testing.T) {
-				if err := runDiff(ctx, test, update, headers); err != nil {
+				if err := runDiff(ctx, test, update); err != nil {
 					t.Fatal(err)
 				}
 			})
@@ -212,6 +217,7 @@
 	name                      string
 	tasks                     chromedp.Tasks
 	urlA, urlB                string
+	headers                   map[string]interface{}
 	cacheA, cacheB            bool
 	outImgA, outImgB, outDiff string
 	viewportWidth             int
@@ -225,17 +231,21 @@
 }
 
 // readTests parses the testcases from a text file.
-func readTests(file string) ([]*testcase, error) {
-	f, err := os.Open(file)
+func readTests(file string, vars map[string]string) ([]*testcase, error) {
+	tmpl, err := template.ParseFiles(file)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("template.ParseFiles(%q): %w", file, err)
 	}
-	defer f.Close()
+	var tmplout bytes.Buffer
+	if err := tmpl.Execute(&tmplout, vars); err != nil {
+		return nil, fmt.Errorf("tmpl.Execute(...): %w", err)
+	}
 	var tests []*testcase
 	var (
 		testName, pathname string
 		tasks              chromedp.Tasks
 		originA, originB   string
+		headers            map[string]interface{}
 		cacheA, cacheB     bool
 		width, height      int
 		lineNo             int
@@ -249,7 +259,7 @@
 	if err != nil {
 		return nil, fmt.Errorf("outDir(%q, %q): %w", dir, file, err)
 	}
-	scan := bufio.NewScanner(f)
+	scan := bufio.NewScanner(&tmplout)
 	for scan.Scan() {
 		lineNo += 1
 		line := strings.TrimSpace(scan.Text())
@@ -269,6 +279,9 @@
 			origins := strings.Split(args, " ")
 			originA, originB = origins[0], origins[1]
 			cacheA, cacheB = false, false
+			if headers != nil {
+				headers = make(map[string]interface{})
+			}
 			if strings.HasSuffix(originA, cacheSuffix) {
 				originA = strings.TrimSuffix(originA, cacheSuffix)
 				cacheA = true
@@ -283,6 +296,15 @@
 			if _, err := url.Parse(originB); err != nil {
 				return nil, fmt.Errorf("url.Parse(%q): %w", originB, err)
 			}
+		case "HEADER":
+			if headers == nil {
+				headers = make(map[string]interface{})
+			}
+			parts := strings.SplitN(args, ":", 2)
+			if len(parts) != 2 {
+				log.Fatalf("invalid header %s on line %d", args, lineNo)
+			}
+			headers[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
 		case "OUTPUT":
 			out, err = outDir(args, "")
 			if err != nil {
@@ -340,10 +362,11 @@
 				return nil, fmt.Errorf("url.Parse(%q): %w", originB+pathname, err)
 			}
 			test := &testcase{
-				name:  testName,
-				tasks: tasks,
-				urlA:  urlA.String(),
-				urlB:  urlB.String(),
+				name:    testName,
+				tasks:   tasks,
+				urlA:    urlA.String(),
+				urlB:    urlB.String(),
+				headers: headers,
 				// Default to viewportScreenshot
 				screenshotType: viewportScreenshot,
 				viewportWidth:  width,
@@ -429,13 +452,13 @@
 
 // runDiff generates screenshots for a given test case and
 // a diff if the screenshots do not match.
-func runDiff(ctx context.Context, test *testcase, update bool, headers map[string]interface{}) error {
+func runDiff(ctx context.Context, test *testcase, update bool) error {
 	fmt.Printf("test %s\n", test.name)
-	screenA, err := screenshot(ctx, test, test.urlA, test.outImgA, test.cacheA, update, headers)
+	screenA, err := screenshot(ctx, test, test.urlA, test.outImgA, test.cacheA, update)
 	if err != nil {
 		return fmt.Errorf("screenshot(ctx, %q, %q, %q, %v): %w", test, test.urlA, test.outImgA, test.cacheA, err)
 	}
-	screenB, err := screenshot(ctx, test, test.urlB, test.outImgB, test.cacheB, update, headers)
+	screenB, err := screenshot(ctx, test, test.urlB, test.outImgB, test.cacheB, update)
 	if err != nil {
 		return fmt.Errorf("screenshot(ctx, %q, %q, %q, %v): %w", test, test.urlB, test.outImgB, test.cacheB, err)
 	}
@@ -474,7 +497,7 @@
 // attempt to read the screenshot from a cache or capture a new screenshot
 // and write it to the cache if it does not exist.
 func screenshot(ctx context.Context, test *testcase,
-	url, file string, cache, update bool, headers map[string]interface{},
+	url, file string, cache, update bool,
 ) (_ *image.Image, err error) {
 	var data []byte
 	// If cache is enabled, try to read the file from the cache.
@@ -491,7 +514,7 @@
 	// If cache is false, this is the first test run, or an update is requested
 	// we capture a new screenshot from a live URL.
 	if !cache || update {
-		data, err = captureScreenshot(ctx, test, url, headers)
+		data, err = captureScreenshot(ctx, test, url)
 		if err != nil {
 			return nil, fmt.Errorf("captureScreenshot(ctx, %q, %q): %w", url, test, err)
 		}
@@ -512,17 +535,15 @@
 
 // captureScreenshot runs a series of browser actions and takes a screenshot
 // of the resulting webpage in an instance of headless chrome.
-func captureScreenshot(ctx context.Context, test *testcase,
-	url string, headers map[string]interface{},
-) ([]byte, error) {
+func captureScreenshot(ctx context.Context, test *testcase, url string) ([]byte, error) {
 	var buf []byte
 	ctx, cancel := chromedp.NewContext(ctx)
 	defer cancel()
 	ctx, cancel = context.WithTimeout(ctx, time.Minute)
 	defer cancel()
 	var tasks chromedp.Tasks
-	if headers != nil {
-		tasks = append(tasks, network.SetExtraHTTPHeaders(headers))
+	if test.headers != nil {
+		tasks = append(tasks, network.SetExtraHTTPHeaders(test.headers))
 	}
 	tasks = append(tasks,
 		chromedp.EmulateViewport(int64(test.viewportWidth), int64(test.viewportHeight)),
diff --git a/internal/screentest/screentest_test.go b/internal/screentest/screentest_test.go
index 86a35d8..8d68f0c 100644
--- a/internal/screentest/screentest_test.go
+++ b/internal/screentest/screentest_test.go
@@ -119,6 +119,7 @@
 					urlA:           "https://pkg.go.dev/about",
 					cacheA:         true,
 					urlB:           "http://localhost:8080/about",
+					headers:        map[string]interface{}{"Authorization": "Bearer token"},
 					outImgA:        filepath.Join(cache, "readtests-txt", "about.pkg-go-dev.png"),
 					outImgB:        filepath.Join(cache, "readtests-txt", "about.localhost-8080.png"),
 					outDiff:        filepath.Join(cache, "readtests-txt", "about.diff.png"),
@@ -131,6 +132,7 @@
 					urlA:           "https://pkg.go.dev/eval",
 					cacheA:         true,
 					urlB:           "http://localhost:8080/eval",
+					headers:        map[string]interface{}{"Authorization": "Bearer token"},
 					outImgA:        filepath.Join(cache, "readtests-txt", "eval.pkg-go-dev.png"),
 					outImgB:        filepath.Join(cache, "readtests-txt", "eval.localhost-8080.png"),
 					outDiff:        filepath.Join(cache, "readtests-txt", "eval.diff.png"),
@@ -147,7 +149,7 @@
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := readTests(tt.args.filename)
+			got, err := readTests(tt.args.filename, map[string]string{"Authorization": "Bearer token"})
 			if (err != nil) != tt.wantErr {
 				t.Errorf("readTests() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -259,6 +261,7 @@
 		urlA:              "http://localhost:6061",
 		cacheA:            true,
 		urlB:              "http://localhost:6061",
+		headers:           map[string]interface{}{"Authorization": "Bearer token"},
 		outImgA:           filepath.Join("testdata", "screenshots", "headers", "headers-test.localhost-6061.png"),
 		outImgB:           filepath.Join("testdata", "screenshots", "headers", "headers-test.localhost-6061.png"),
 		outDiff:           filepath.Join("testdata", "screenshots", "headers", "headers-test.diff.png"),
@@ -266,7 +269,7 @@
 		viewportHeight:    960,
 		screenshotType:    elementScreenshot,
 		screenshotElement: "#result",
-	}, false, map[string]interface{}{"Authorization": "Bearer token"}); err != nil {
+	}, false); err != nil {
 		t.Fatal(err)
 	}
 }
diff --git a/internal/screentest/testdata/readtests.txt b/internal/screentest/testdata/readtests.txt
index 082377a..d1f1f7a 100644
--- a/internal/screentest/testdata/readtests.txt
+++ b/internal/screentest/testdata/readtests.txt
@@ -24,6 +24,7 @@
 capture viewport 540x1080
 
 compare https://pkg.go.dev::cache http://localhost:8080
+header Authorization: {{.Authorization}}
 
 pathname /about
 capture