Repository: google/exposure-notifications-server
Pull request #1344 (2021-02-04 20:28:27)

Refactor generate service to be a proper server, use new middlewares
**Release Note**

<!-- Enter your extended release note in the below block. If the PR requires
additional action from users switching to the new release, include the string
"action required". If no release note is required, write "NONE". -->

```release-note
Refactor generate service to be a proper server, use new middlewares
```

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1344#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sethvargo: /assign @whaught 
google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1344#pullrequestreview-583787800" title="Approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/1344#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [mikehelmick,sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/cmd/generate/main.go b/cmd/generate/main.go
index 9eb3cad..669fe27 100644
--- a/cmd/generate/main.go
+++ b/cmd/generate/main.go
@@ -26,7 +26,6 @@ import (
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	_ "github.com/google/exposure-notifications-server/pkg/observability"
 	"github.com/google/exposure-notifications-server/pkg/server"
-	"github.com/gorilla/mux"
 	"github.com/sethvargo/go-signalcontext"
 )
 
@@ -57,27 +56,22 @@ func main() {
 func realMain(ctx context.Context) error {
 	logger := logging.FromContext(ctx)
 
-	var config generate.Config
-	env, err := setup.Setup(ctx, &config)
+	var cfg generate.Config
+	env, err := setup.Setup(ctx, &cfg)
 	if err != nil {
 		return fmt.Errorf("setup.Setup: %w", err)
 	}
 	defer env.Close(ctx)
 
-	handler, err := generate.NewHandler(ctx, &config, env)
+	generateServer, err := generate.NewServer(&cfg, env)
 	if err != nil {
-		return fmt.Errorf("generate.NewHandler: %w", err)
+		return fmt.Errorf("generate.NewServer: %w", err)
 	}
 
-	r := mux.NewRouter()
-	r.Handle("/", handler)
-	r.Handle("/health", server.HandleHealthz(ctx))
-
-	srv, err := server.New(config.Port)
+	srv, err := server.New(cfg.Port)
 	if err != nil {
 		return fmt.Errorf("server.New: %w", err)
 	}
-	logger.Infof("listening on :%s", config.Port)
-
-	return srv.ServeHTTPHandler(ctx, r)
+	logger.Infow("server listening", "port", cfg.Port)
+	return srv.ServeHTTPHandler(ctx, generateServer.Routes(ctx))
 }
diff --git a/cmd/jwks/main.go b/cmd/jwks/main.go
index a14b7aa..a906408 100644
--- a/cmd/jwks/main.go
+++ b/cmd/jwks/main.go
@@ -55,23 +55,22 @@ func main() {
 func realMain(ctx context.Context) error {
 	logger := logging.FromContext(ctx)
 
-	var config jwks.Config
-	env, err := setup.Setup(ctx, &config)
+	var cfg jwks.Config
+	env, err := setup.Setup(ctx, &cfg)
 	if err != nil {
 		return fmt.Errorf("setup.Setup: %w", err)
 	}
 	defer env.Close(ctx)
 
-	jwksServer, err := jwks.NewServer(&config, env)
+	jwksServer, err := jwks.NewServer(&cfg, env)
 	if err != nil {
 		return fmt.Errorf("jwks.NewServer: %w", err)
 	}
 
-	srv, err := server.New(config.Port)
+	srv, err := server.New(cfg.Port)
 	if err != nil {
 		return fmt.Errorf("server.New: %w", err)
 	}
-	logger.Infof("listening on :%s", config.Port)
-
+	logger.Infow("server listening", "port", cfg.Port)
 	return srv.ServeHTTPHandler(ctx, jwksServer.Routes(ctx))
 }
diff --git a/internal/generate/handler.go b/internal/generate/handle_generate.go
similarity index 63%
rename from internal/generate/handler.go
rename to internal/generate/handle_generate.go
index cfa532f..0de2820 100644
--- a/internal/generate/handler.go
+++ b/internal/generate/handle_generate.go
@@ -1,4 +1,4 @@
-// Copyright 2020 Google LLC
+// Copyright 2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -22,119 +22,87 @@ import (
 	"strings"
 	"time"
 
-	"go.opencensus.io/trace"
-	"go.uber.org/zap"
-
 	"github.com/google/exposure-notifications-server/internal/pb"
 	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
-	"github.com/google/exposure-notifications-server/internal/serverenv"
 	"github.com/google/exposure-notifications-server/internal/verification"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/google/exposure-notifications-server/pkg/timeutils"
 	"github.com/google/exposure-notifications-server/pkg/util"
+	"go.opencensus.io/trace"
 )
 
-func NewHandler(ctx context.Context, config *Config, env *serverenv.ServerEnv) (http.Handler, error) {
-	logger := logging.FromContext(ctx).Named("generate")
-
-	if env.Database() == nil {
-		return nil, fmt.Errorf("missing database in server environment")
-	}
-
-	transformer, err := publishmodel.NewTransformer(config)
-	if err != nil {
-		return nil, fmt.Errorf("model.NewTransformer: %w", err)
-	}
-	logger.Debugw("max keys per upload", "value", config.MaxKeysOnPublish)
-	logger.Debugw("max interval start age", "value", config.MaxIntervalAge)
-	logger.Debugw("truncate window", "value", config.TruncateWindow)
-
-	return &generateHandler{
-		serverenv:   env,
-		transformer: transformer,
-		config:      config,
-		database:    publishdb.New(env.Database()),
-		logger:      logger,
-	}, nil
-}
-
-type generateHandler struct {
-	config      *Config
-	serverenv   *serverenv.ServerEnv
-	transformer *publishmodel.Transformer
-	database    *publishdb.PublishDB
-	logger      *zap.SugaredLogger
-}
-
-func (h *generateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	ctx, span := trace.StartSpan(r.Context(), "(*generate.generateHandler).ServeHTTP")
-	defer span.End()
+func (s *Server) handleGenerate() http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		ctx, span := trace.StartSpan(r.Context(), "(*generate.Server).handleGenerate")
+		defer span.End()
 
-	logger := h.logger.Named("ServeHTTP")
+		logger := logging.FromContext(ctx).Named("handleGenerate")
 
-	regionStr := h.config.DefaultRegion
-	if v := r.URL.Query().Get("region"); v != "" {
-		regionStr = v
-	}
-	regions := strings.Split(regionStr, ",")
+		regionStr := s.config.DefaultRegion
+		if v := r.URL.Query().Get("region"); v != "" {
+			regionStr = v
+		}
+		regions := strings.Split(regionStr, ",")
 
-	logger.Debugw("generating data", "regions", regions)
+		logger.Debugw("generating data", "regions", regions)
 
-	if err := h.generate(ctx, regions); err != nil {
-		logger.Errorw("failed to generate", "error", err)
-		span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: err.Error()})
-		w.WriteHeader(http.StatusInternalServerError)
-		fmt.Fprint(w, err.Error())
-		return
-	}
+		if err := s.generate(ctx, regions); err != nil {
+			logger.Errorw("failed to generate", "error", err)
+			span.SetStatus(trace.Status{Code: trace.StatusCodeInternal, Message: err.Error()})
+			w.WriteHeader(http.StatusInternalServerError)
+			fmt.Fprint(w, err.Error())
+			return
+		}
 
-	w.WriteHeader(http.StatusOK)
-	fmt.Fprint(w, "successfully generated exposure keys")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, "successfully generated exposure keys")
+	})
 }
 
-func (h *generateHandler) generate(ctx context.Context, regions []string) error {
+func (s *Server) generate(ctx context.Context, regions []string) error {
 	for _, r := range regions {
-		if err := h.generateKeysInRegion(ctx, r); err != nil {
+		if err := s.generateKeysInRegion(ctx, r); err != nil {
 			return fmt.Errorf("generateKeysInRegion: %w", err)
 		}
 	}
 	return nil
 }
 
-func (h *generateHandler) generateKeysInRegion(ctx context.Context, region string) error {
+func (s *Server) generateKeysInRegion(ctx context.Context, region string) error {
+	logger := logging.FromContext(ctx).Named("generateKeysInRegion")
+
 	// We require at least 2 keys because revision only revises a subset of keys,
 	// and that subset selects a random sample from (0-len(keys)], and rand panics
 	// if you try to generate a random number between 0 and 0 :).
-	if h.config.KeysPerExposure < 2 {
+	if s.config.KeysPerExposure < 2 {
 		return fmt.Errorf("number of keys to publish must be at least 2")
 	}
 
-	logger := h.logger.Named("generate")
 	// API calls treat region as a list, for legacy regions.
 	regions := []string{region}
 
 	traveler := false
 	if val, err := util.RandomInt(100); err != nil {
 		return fmt.Errorf("failed to determien traveler status: %w", err)
-	} else if val < h.config.ChanceOfTraveler {
+	} else if val < s.config.ChanceOfTraveler {
 		traveler = true
 	}
 
 	now := time.Now().UTC()
 	// Find the valid intervals - starting with today and working backwards
-	minInterval := publishmodel.IntervalNumber(timeutils.UTCMidnight(now.Add(-1 * h.config.MaxIntervalAge).Add(24 * time.Hour)))
+	minInterval := publishmodel.IntervalNumber(timeutils.UTCMidnight(now.Add(-1 * s.config.MaxIntervalAge).Add(24 * time.Hour)))
 	curInterval := publishmodel.IntervalNumber(timeutils.UTCMidnight(now))
-	intervals := make([]int32, 0, h.config.KeysPerExposure)
-	for i := 0; i < h.config.KeysPerExposure && curInterval >= minInterval; i++ {
+	intervals := make([]int32, 0, s.config.KeysPerExposure)
+	for i := 0; i < s.config.KeysPerExposure && curInterval >= minInterval; i++ {
 		intervals = append(intervals, curInterval)
 		curInterval -= verifyapi.MaxIntervalCount
 	}
 
 	batchTime := now
-	for i := 0; i < h.config.NumExposures; i++ {
-		logger.Debugf("generating exposure %d of %d", i+1, h.config.NumExposures)
+	for i := 0; i < s.config.NumExposures; i++ {
+		logger.Debugf("generating exposure %d of %d", i+1, s.config.NumExposures)
 
 		exposures, err := util.GenerateExposuresForIntervals(intervals)
 		if err != nil {
@@ -146,7 +114,7 @@ func (h *generateHandler) generateKeysInRegion(ctx context.Context, region strin
 			Traveler:          traveler,
 		}
 
-		if h.config.SimulateSameDayRelease {
+		if s.config.SimulateSameDayRelease {
 			sort.Slice(publish.Keys, func(i int, j int) bool {
 				return publish.Keys[i].IntervalNumber < publish.Keys[j].IntervalNumber
 			})
@@ -165,7 +133,7 @@ func (h *generateHandler) generateKeysInRegion(ctx context.Context, region strin
 		if err != nil {
 			return fmt.Errorf("failed to decide revised key status: %w", err)
 		}
-		generateRevisedKeys := val <= h.config.ChanceOfKeyRevision
+		generateRevisedKeys := val <= s.config.ChanceOfKeyRevision
 
 		reportType := verifyapi.ReportTypeClinical
 		if !generateRevisedKeys {
@@ -185,12 +153,12 @@ func (h *generateHandler) generateKeysInRegion(ctx context.Context, region strin
 			SymptomOnsetInterval: uint32(publish.Keys[intervalIdx].IntervalNumber),
 		}
 
-		result, err := h.transformer.TransformPublish(ctx, publish, regions, &claims, batchTime)
+		result, err := s.transformer.TransformPublish(ctx, publish, regions, &claims, batchTime)
 		if err != nil {
 			return fmt.Errorf("failed to transform generated exposures: %w", err)
 		}
 
-		n, err := h.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+		n, err := s.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
 			Incoming:     result.Exposures,
 			RequireToken: true,
 		})
@@ -206,9 +174,9 @@ func (h *generateHandler) generateKeysInRegion(ctx context.Context, region strin
 			}
 
 			claims.ReportType = revisedReportType
-			batchTime = batchTime.Add(h.config.KeyRevisionDelay)
+			batchTime = batchTime.Add(s.config.KeyRevisionDelay)
 
-			result, err := h.transformer.TransformPublish(ctx, publish, regions, &claims, batchTime)
+			result, err := s.transformer.TransformPublish(ctx, publish, regions, &claims, batchTime)
 			if err != nil {
 				return fmt.Errorf("failed to transform generated exposures: %w", err)
 			}
@@ -224,7 +192,7 @@ func (h *generateHandler) generateKeysInRegion(ctx context.Context, region strin
 			}
 
 			// Bypass revision token enforcement on generated data.
-			n, err := h.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
+			n, err := s.database.InsertAndReviseExposures(ctx, &publishdb.InsertAndReviseExposuresRequest{
 				Incoming:     result.Exposures,
 				Token:        &token,
 				RequireToken: true,
diff --git a/internal/generate/server.go b/internal/generate/server.go
new file mode 100644
index 0000000..ccf4c9c
--- /dev/null
+++ b/internal/generate/server.go
@@ -0,0 +1,67 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package generate
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/exposure-notifications-server/internal/middleware"
+	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
+	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
+	"github.com/google/exposure-notifications-server/internal/serverenv"
+	"github.com/google/exposure-notifications-server/pkg/logging"
+	"github.com/google/exposure-notifications-server/pkg/server"
+	"github.com/gorilla/mux"
+)
+
+type Server struct {
+	config      *Config
+	serverenv   *serverenv.ServerEnv
+	transformer *publishmodel.Transformer
+	database    *publishdb.PublishDB
+}
+
+func NewServer(cfg *Config, env *serverenv.ServerEnv) (*Server, error) {
+	if env.Database() == nil {
+		return nil, fmt.Errorf("missing database in server environment")
+	}
+
+	transformer, err := publishmodel.NewTransformer(cfg)
+	if err != nil {
+		return nil, fmt.Errorf("model.NewTransformer: %w", err)
+	}
+
+	return &Server{
+		serverenv:   env,
+		transformer: transformer,
+		config:      cfg,
+		database:    publishdb.New(env.Database()),
+	}, nil
+}
+
+// Routes returns the router for this server.
+func (s *Server) Routes(ctx context.Context) *mux.Router {
+	logger := logging.FromContext(ctx).Named("generate")
+
+	r := mux.NewRouter()
+	r.Use(middleware.PopulateRequestID())
+	r.Use(middleware.PopulateLogger(logger))
+
+	r.Handle("/health", server.HandleHealthz(ctx))
+	r.Handle("/", s.handleGenerate())
+
+	return r
+}
diff --git a/internal/project/context.go b/internal/project/context.go
new file mode 100644
index 0000000..e1bbfbb
--- /dev/null
+++ b/internal/project/context.go
@@ -0,0 +1,41 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package project
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/exposure-notifications-server/pkg/logging"
+
+	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest"
+)
+
+// TestContext returns a context with test values pre-populated.
+func TestContext(tb testing.TB) context.Context {
+	ctx := context.Background()
+	ctx = logging.WithLogger(ctx, TestLogger(tb))
+	return ctx
+}
+
+// TestLogger returns a logger configured for test. See the following link for
+// more information:
+//
+//     https://pkg.go.dev/go.uber.org/zap/zaptest
+//
+func TestLogger(tb testing.TB) *zap.SugaredLogger {
+	return zaptest.NewLogger(tb, zaptest.Level(zap.WarnLevel)).Sugar()
+}
diff --git a/pkg/logging/logger.go b/pkg/logging/logger.go
index a1941ef..c1ff167 100644
--- a/pkg/logging/logger.go
+++ b/pkg/logging/logger.go
@@ -22,7 +22,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/google/exposure-notifications-server/internal/project"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -75,7 +74,7 @@ func NewLogger(level string, development bool) *zap.SugaredLogger {
 // parameters.
 func NewLoggerFromEnv() *zap.SugaredLogger {
 	level := os.Getenv("LOG_LEVEL")
-	development := strings.ToLower(project.TrimSpaceAndNonPrintable(os.Getenv("LOG_MODE"))) == "development"
+	development := strings.ToLower(strings.TrimSpace(os.Getenv("LOG_MODE"))) == "development"
 	return NewLogger(level, development)
 }
 
@@ -155,7 +154,7 @@ var developmentEncoderConfig = zapcore.EncoderConfig{
 // levelToZapLevel converts the given string to the appropriate zap level
 // value.
 func levelToZapLevel(s string) zapcore.Level {
-	switch strings.ToUpper(project.TrimSpaceAndNonPrintable(s)) {
+	switch strings.ToUpper(strings.TrimSpace(s)) {
 	case levelDebug:
 		return zapcore.DebugLevel
 	case levelInfo:


Test patch:
diff --git a/internal/generate/generate_test.go b/internal/generate/generate_test.go
index 2bde8c4..0278d3e 100644
--- a/internal/generate/generate_test.go
+++ b/internal/generate/generate_test.go
@@ -15,7 +15,6 @@
 package generate
 
 import (
-	"context"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -23,9 +22,9 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/project"
 	publishdb "github.com/google/exposure-notifications-server/internal/publish/database"
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
-	"github.com/google/exposure-notifications-server/internal/serverenv"
 	"github.com/google/exposure-notifications-server/internal/setup"
 	"github.com/google/exposure-notifications-server/pkg/observability"
 	"github.com/google/exposure-notifications-server/pkg/secrets"
@@ -40,13 +39,13 @@ func TestMain(m *testing.M) {
 	m.Run()
 }
 
-func testHandler(tb testing.TB) (*Config, *serverenv.ServerEnv) {
+func testServer(tb testing.TB) *Server {
 	tb.Helper()
 
-	ctx := context.Background()
+	ctx := project.TestContext(tb)
 	_, dbConfig := testDatabaseInstance.NewDatabase(tb)
 
-	config := &Config{
+	cfg := &Config{
 		Database: *dbConfig,
 		SecretManager: secrets.Config{
 			SecretManagerType: secrets.SecretManagerTypeInMemory,
@@ -67,7 +66,7 @@ func testHandler(tb testing.TB) (*Config, *serverenv.ServerEnv) {
 		KeyRevisionDelay:             1 * time.Hour,
 	}
 
-	env, err := setup.SetupWith(ctx, config, envconfig.MapLookuper(map[string]string{
+	env, err := setup.SetupWith(ctx, cfg, envconfig.MapLookuper(map[string]string{
 		// We want this to be 0, but envconfig overrides the default value of 0 when
 		// processing, so force it to be 0 by simulating the environment variable
 		// value to 0.
@@ -82,43 +81,48 @@ func testHandler(tb testing.TB) (*Config, *serverenv.ServerEnv) {
 		}
 	})
 
-	return config, env
+	srv, err := NewServer(cfg, env)
+	if err != nil {
+		tb.Fatal(err)
+	}
+	return srv
 }
 
-func TestGenerateHandler_ServeHTTP(t *testing.T) {
+func TestServer_smoke(t *testing.T) {
 	// Note: this only tests that the server is capable of sending and receiving
 	// responses. It does not attempt to exhaustively exercise the actual
 	// generation, which is tested in a different function.
 	t.Parallel()
 
-	ctx := context.Background()
-	config, env := testHandler(t)
-	h, err := NewHandler(ctx, config, env)
-	if err != nil {
-		t.Fatal(err)
-	}
+	ctx := project.TestContext(t)
 
-	req, err := http.NewRequest("POST", "/", nil)
-	if err != nil {
-		t.Fatal(err)
-	}
+	srv := testServer(t)
+	mux := srv.Routes(ctx)
 
-	rr := httptest.NewRecorder()
-	handler := http.HandlerFunc(h.ServeHTTP)
-	handler.ServeHTTP(rr, req)
+	t.Run("handleGenerate", func(t *testing.T) {
+		t.Parallel()
 
-	if got, want := rr.Code, http.StatusOK; got != want {
-		t.Errorf("expected %#v to be %#v", got, want)
-	}
+		r := httptest.NewRequest("POST", "/", nil)
+		w := httptest.NewRecorder()
+		mux.ServeHTTP(w, r)
 
-	if got, want := rr.Body.String(), "successfully generated exposure keys"; got != want {
-		t.Errorf("expected %#v to be %#v", got, want)
-	}
+		w.Flush()
+
+		if got, want := w.Code, http.StatusOK; got != want {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+
+		if got, want := w.Body.String(), "successfully generated exposure keys"; got != want {
+			t.Errorf("expected %#v to be %#v", got, want)
+		}
+	})
 }
 
-func TestGenerateHandler_Generate(t *testing.T) {
+func TestServer_handleGenerate(t *testing.T) {
 	t.Parallel()
 
+	ctx := project.TestContext(t)
+
 	cases := []struct {
 		name string
 		err  string
@@ -149,27 +153,16 @@ func TestGenerateHandler_Generate(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
-			config, env := testHandler(t)
-			db := env.Database()
+			srv := testServer(t)
 
 			if tc.sameDayRelease {
-				config.SimulateSameDayRelease = true
+				srv.config.SimulateSameDayRelease = true
 			}
 			if tc.reviseKeys {
-				config.ChanceOfKeyRevision = 100
-			}
-
-			ctx := context.Background()
-			h, err := NewHandler(ctx, config, env)
-			if err != nil {
-				t.Fatal(err)
-			}
-			gh, ok := h.(*generateHandler)
-			if !ok {
-				t.Fatalf("expected handler to be generateHandler")
+				srv.config.ChanceOfKeyRevision = 100
 			}
 
-			if err := gh.generate(ctx, tc.regions); err != nil {
+			if err := srv.generate(ctx, tc.regions); err != nil {
 				if tc.err == "" {
 					t.Fatal(err)
 				}
@@ -179,7 +172,7 @@ func TestGenerateHandler_Generate(t *testing.T) {
 			}
 
 			var exposures []*publishmodel.Exposure
-			if _, err := publishdb.New(db).IterateExposures(ctx, publishdb.IterateExposuresCriteria{}, func(m *publishmodel.Exposure) error {
+			if _, err := srv.database.IterateExposures(ctx, publishdb.IterateExposuresCriteria{}, func(m *publishmodel.Exposure) error {
 				exposures = append(exposures, m)
 				return nil
 			}); err != nil {
diff --git a/internal/jwks/jwks_test.go b/internal/jwks/jwks_test.go
index 8194e77..f6771ba 100644
--- a/internal/jwks/jwks_test.go
+++ b/internal/jwks/jwks_test.go
@@ -15,7 +15,6 @@
 package jwks
 
 import (
-	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -25,6 +24,7 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/internal/project"
 	hadb "github.com/google/exposure-notifications-server/internal/verification/database"
 	"github.com/google/exposure-notifications-server/internal/verification/model"
 )
@@ -137,7 +137,7 @@ func TestUpdateHA(t *testing.T) {
 			defer ts.Close()
 
 			// Set up the test.
-			ctx := context.Background()
+			ctx := project.TestContext(t)
 			testDB, _ := testDatabaseInstance.NewDatabase(t)
 			mgr, err := NewManager(testDB, time.Minute, 5*time.Second)
 			if err != nil {
diff --git a/internal/project/context_test.go b/internal/project/context_test.go
new file mode 100644
index 0000000..0bb3082
--- /dev/null
+++ b/internal/project/context_test.go
@@ -0,0 +1,32 @@
+// Copyright 2021 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package project
+
+import (
+	"testing"
+
+	"github.com/google/exposure-notifications-server/pkg/logging"
+)
+
+func TestTestContext(t *testing.T) {
+	t.Parallel()
+
+	ctx := TestContext(t)
+	logger := logging.FromContext(ctx)
+
+	if logger == nil {
+		t.Error("expected a test logger")
+	}
+}