dashboard/app: allow to set spanner context only from tests

getSpannerClient returns prod client as a default.

9 files changed, 25 insertions, 27 deletions

diff --git a/dashboard/app/api.go b/dashboard/app/api.go
index f491b23a0..52e33dbb0 100644
--- a/dashboard/app/api.go
+++ b/dashboard/app/api.go
@@ -104,7 +104,6 @@ var maxCrashes = func() int {
 func handleJSON(fn JSONHandler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		c := appengine.NewContext(r)
-		c = SetCoverageDBClient(c, coverageDBClient)
 		reply, err := fn(c, r)
 		if err != nil {
 			status := logErrorPrepareStatus(c, err)
@@ -1944,7 +1943,7 @@ func apiSaveCoverage(c context.Context, payload io.Reader) (interface{}, error)
 		sss = service.List()
 		log.Infof(c, "found %d subsystems for %s namespace", len(sss), descr.Namespace)
 	}
-	rowsCreated, err := coveragedb.SaveMergeResult(c, GetCoverageDBClient(c), descr, jsonDec, sss)
+	rowsCreated, err := coveragedb.SaveMergeResult(c, getCoverageDBClient(c), descr, jsonDec, sss)
 	if err != nil {
 		log.Errorf(c, "error storing coverage for ns %s, date %s: %v",
 			descr.Namespace, descr.DateTo.String(), err)
diff --git a/dashboard/app/coverage.go b/dashboard/app/coverage.go
index c600aab9e..e2e0c744f 100644
--- a/dashboard/app/coverage.go
+++ b/dashboard/app/coverage.go
@@ -27,7 +27,7 @@ var coverageDBClient spannerclient.SpannerClient
 func initCoverageDB() {
 	if !appengine.IsAppEngine() {
 		// It is a test environment.
-		// Use SetCoverageDBClient to specify the coveragedb mock or emulator in every test.
+		// Use setCoverageDBClient to specify the coveragedb mock or emulator in every test.
 		return
 	}
 	projectID := os.Getenv("GOOGLE_CLOUD_PROJECT")
@@ -40,13 +40,15 @@ func initCoverageDB() {
 
 var keyCoverageDBClient = "coveragedb client key"
 
-func SetCoverageDBClient(ctx context.Context, client spannerclient.SpannerClient) context.Context {
-	return context.WithValue(ctx, &keyCoverageDBClient, client)
-}
-
-func GetCoverageDBClient(ctx context.Context) spannerclient.SpannerClient {
-	client, _ := ctx.Value(&keyCoverageDBClient).(spannerclient.SpannerClient)
-	return client
+func getCoverageDBClient(ctx context.Context) spannerclient.SpannerClient {
+	ctxClient, _ := ctx.Value(&keyCoverageDBClient).(spannerclient.SpannerClient)
+	if ctxClient == nil && coverageDBClient == nil {
+		panic("attempt to get coverage db client before it was set in tests")
+	}
+	if ctxClient != nil {
+		return ctxClient
+	}
+	return coverageDBClient
 }
 
 type funcStyleBodyJS func(
@@ -190,7 +192,7 @@ func handleHeatmap(c context.Context, w http.ResponseWriter, r *http.Request, f
 
 	var style template.CSS
 	var body, js template.HTML
-	if style, body, js, err = f(c, GetCoverageDBClient(c),
+	if style, body, js, err = f(c, getCoverageDBClient(c),
 		&coveragedb.SelectScope{
 			Ns:        hdr.Namespace,
 			Subsystem: ss,
@@ -261,7 +263,7 @@ func handleFileCoverage(c context.Context, w http.ResponseWriter, r *http.Reques
 	}
 	onlyUnique := r.FormValue(covPageParams[UniqueOnly]) == "1"
 	mainNsRepo, _ := nsConfig.mainRepoBranch()
-	client := GetCoverageDBClient(c)
+	client := getCoverageDBClient(c)
 	if client == nil {
 		return fmt.Errorf("spannerdb client is nil")
 	}
@@ -329,7 +331,7 @@ func handleCoverageGraph(c context.Context, w http.ResponseWriter, r *http.Reque
 	if periodType != coveragedb.QuarterPeriod && periodType != coveragedb.MonthPeriod {
 		return fmt.Errorf("only quarter and month are allowed, but received %s instead", periodType)
 	}
-	hist, err := MergedCoverage(c, GetCoverageDBClient(c), hdr.Namespace, periodType)
+	hist, err := MergedCoverage(c, getCoverageDBClient(c), hdr.Namespace, periodType)
 	if err != nil {
 		return err
 	}
diff --git a/dashboard/app/coverage_test.go b/dashboard/app/coverage_test.go
index 516409228..bdd528e43 100644
--- a/dashboard/app/coverage_test.go
+++ b/dashboard/app/coverage_test.go
@@ -4,6 +4,7 @@
 package main
 
 import (
+	"context"
 	"strings"
 	"testing"
 
@@ -17,6 +18,10 @@ import (
 	"google.golang.org/api/iterator"
 )
 
+func setCoverageDBClient(ctx context.Context, client spannerclient.SpannerClient) context.Context {
+	return context.WithValue(ctx, &keyCoverageDBClient, client)
+}
+
 func TestFileCoverage(t *testing.T) {
 	tests := []struct {
 		name      string
diff --git a/dashboard/app/handler.go b/dashboard/app/handler.go
index 77a22e4c0..157e33383 100644
--- a/dashboard/app/handler.go
+++ b/dashboard/app/handler.go
@@ -18,7 +18,6 @@ import (
 	"time"
 
 	"github.com/google/syzkaller/pkg/html"
-	"google.golang.org/appengine/v2"
 	"google.golang.org/appengine/v2/log"
 	"google.golang.org/appengine/v2/user"
 )
@@ -33,11 +32,7 @@ func handlerWrapper(fn contextHandler) http.Handler {
 
 func handleContext(fn contextHandler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		c := appengine.NewContext(r)
-		if coverageDBClient != nil { // Nil in prod.
-			c = SetCoverageDBClient(c, coverageDBClient)
-		}
-		c = context.WithValue(c, &currentURLKey, r.URL.RequestURI())
+		c := context.WithValue(r.Context(), &currentURLKey, r.URL.RequestURI())
 		authorizedUser, _ := userAccessLevel(currentUser(c), "", getConfig(c))
 		if !authorizedUser {
 			if !throttleRequest(c, w, r) {
diff --git a/dashboard/app/public_json_api.go b/dashboard/app/public_json_api.go
index 90e2fe8d6..3b882ea24 100644
--- a/dashboard/app/public_json_api.go
+++ b/dashboard/app/public_json_api.go
@@ -188,7 +188,7 @@ func writeExtAPICoverageFor(ctx context.Context, w io.Writer, ns, repo string) e
 		return fmt.Errorf("coveragedb.GenNPeriodsTill: %w", err)
 	}
 	defaultTimePeriod := tps[0]
-	covDBClient := GetCoverageDBClient(ctx)
+	covDBClient := getCoverageDBClient(ctx)
 	ff, err := coveragedb.MakeFuncFinder(ctx, covDBClient, ns, defaultTimePeriod)
 	if err != nil {
 		return fmt.Errorf("coveragedb.MakeFuncFinder: %w", err)
diff --git a/dashboard/app/public_json_api_test.go b/dashboard/app/public_json_api_test.go
index f0aee9f79..5a23347db 100644
--- a/dashboard/app/public_json_api_test.go
+++ b/dashboard/app/public_json_api_test.go
@@ -259,7 +259,7 @@ func TestPublicJSONAPI(t *testing.T) {
 }
 
 func TestWriteExtAPICoverageFor(t *testing.T) {
-	ctx := SetCoverageDBClient(context.Background(), fileFuncLinesDBFixture(t,
+	ctx := setCoverageDBClient(context.Background(), fileFuncLinesDBFixture(t,
 		[]*coveragedb.FuncLines{
 			{
 				FilePath: "/file",
diff --git a/dashboard/app/reporting_email.go b/dashboard/app/reporting_email.go
index a141caba5..f15de3287 100644
--- a/dashboard/app/reporting_email.go
+++ b/dashboard/app/reporting_email.go
@@ -110,9 +110,6 @@ func (cfg *EmailConfig) Validate() error {
 // Assuming it is called June 15, the monthly report will cover April-May diff.
 func handleCoverageReports(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
-	if coverageDBClient != nil { // initialized in prod deployment, nil in tests
-		ctx = SetCoverageDBClient(r.Context(), coverageDBClient)
-	}
 	targetDate := civil.DateOf(timeNow(ctx)).AddMonths(-1)
 	periods, err := coveragedb.GenNPeriodsTill(2, targetDate, "month")
 	if err != nil {
@@ -196,7 +193,7 @@ func sendNsCoverageReport(ctx context.Context, ns, email string,
 func coverageTable(ctx context.Context, ns string, fromTo []coveragedb.TimePeriod, minDrop int) (string, error) {
 	covAndDates, err := coveragedb.FilesCoverageWithDetails(
 		ctx,
-		GetCoverageDBClient(ctx),
+		getCoverageDBClient(ctx),
 		&coveragedb.SelectScope{
 			Ns:      ns,
 			Periods: fromTo,
diff --git a/dashboard/app/reporting_test.go b/dashboard/app/reporting_test.go
index 44a040f11..5aa9368b9 100644
--- a/dashboard/app/reporting_test.go
+++ b/dashboard/app/reporting_test.go
@@ -1399,7 +1399,7 @@ func TestCoverageRegression(t *testing.T) {
 		Return(mTran2).Once()
 
 	c.transformContext = func(ctx context.Context) context.Context {
-		return SetCoverageDBClient(ctx, m)
+		return setCoverageDBClient(ctx, m)
 	}
 	_, err := c.AuthGET(AccessAdmin, "/cron/email_coverage_reports")
 	assert.NoError(t, err)
diff --git a/dashboard/app/util_test.go b/dashboard/app/util_test.go
index be7841b27..014e34a57 100644
--- a/dashboard/app/util_test.go
+++ b/dashboard/app/util_test.go
@@ -236,7 +236,7 @@ func (c *Ctx) setCoverageMocks(ns string, dbClientMock spannerclient.SpannerClie
 			ret.Coverage = &CoverageConfig{WebGitURI: "test-git"}
 			return &ret
 		})
-		ctxWithSpanner := SetCoverageDBClient(ctx, dbClientMock)
+		ctxWithSpanner := setCoverageDBClient(ctx, dbClientMock)
 		ctxWithSpannerAndFileProvider := setWebGit(ctxWithSpanner, fileProvMock)
 		return contextWithConfig(ctxWithSpannerAndFileProvider, newConfig)
 	}