From ae39f09e82c3c045924298fb8e10e303240b0194 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Mon, 18 May 2020 20:03:28 -0700
Subject: [PATCH 01/12] Publish integration test

---
 testing/integration/publish_test.go | 62 +++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)
 create mode 100644 testing/integration/publish_test.go

diff --git a/testing/integration/publish_test.go b/testing/integration/publish_test.go
new file mode 100644
index 000000000..ede48f6d6
--- /dev/null
+++ b/testing/integration/publish_test.go
@@ -0,0 +1,62 @@
+package integration
+
+import (
+	"log"
+	"testing"
+
+	"github.com/google/exposure-notifications-server/internal/database"
+	"github.com/google/exposure-notifications-server/testing/enclient"
+)
+
+const appPackageName = ""
+
+func TestCleanupAPI(t *testing.T) {
+	var bts []byte
+	requestUrl := "http://localhost:8080/cleanup-exposure"
+	resp, err := enclient.PostRequest(requestUrl, bts)
+	if err != nil {
+		t.Errorf("request failed: %v, %v", err, resp)
+		return
+	}
+
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Cleanup request is sent to %v", requestUrl)
+}
+
+func TestExportAPI(t *testing.T) {
+	var bts []byte
+	requestUrl := "http://localhost:8080/export/create-batches"
+	resp, err := enclient.PostRequest(requestUrl, bts)
+	if err != nil {
+		t.Errorf("request failed: %v, %v", err, resp)
+		return
+	}
+
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Create batches request is sent to %v", requestUrl)
+}
+
+func TestExportWorkerApi(t *testing.T) {
+	var bts []byte
+	requestUrl := "http://localhost:8080/export/do-work"
+	resp, err := enclient.PostRequest(requestUrl, bts)
+	if err != nil {
+		t.Errorf("request failed: %v, %v", err, resp)
+		return
+	}
+
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Export worker request is sent to %v", requestUrl)
+}
+
+func publishKey(keys []database.ExposureKey, regions []string) database.Publish {
+	padding := enclient.RandomBytes(1000)
+	return database.Publish{
+		Keys:                      keys,
+		Regions:                   regions,
+		AppPackageName:            appPackageName,
+		DeviceVerificationPayload: "Test Device Verification Payload",
+		VerificationPayload:       "Test Authority",
+		Padding:                   enclient.ToBase64(padding),
+	}
+}

From 023436b4838478e2c417ed11611b439939235c30 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 20 May 2020 15:20:43 -0700
Subject: [PATCH 02/12] Adding e2e integration test.

---
 go.mod                              |   1 +
 internal/export/exportfile.go       |   1 +
 internal/util/generate.go           |  15 ++++
 testing/enclient/enclient.go        |   2 +-
 testing/enclient/fs_client.go       |  47 ++++++++++
 testing/integration/en_api.go       | 134 ++++++++++++++++++++++++++++
 testing/integration/publish_test.go | 115 +++++++++++++++---------
 7 files changed, 270 insertions(+), 45 deletions(-)
 create mode 100644 testing/enclient/fs_client.go
 create mode 100644 testing/integration/en_api.go

diff --git a/go.mod b/go.mod
index afe42fdd7..2ed222bff 100644
--- a/go.mod
+++ b/go.mod
@@ -54,6 +54,7 @@ require (
 	github.com/oracle/oci-go-sdk v19.3.0+incompatible // indirect
 	github.com/ory/dockertest v3.3.5+incompatible
 	github.com/pierrec/lz4 v2.5.2+incompatible // indirect
+	github.com/pkg/errors v0.9.1
 	github.com/posener/complete v1.2.3 // indirect
 	github.com/prometheus/client_golang v1.6.0 // indirect
 	github.com/prometheus/common v0.10.0 // indirect
diff --git a/internal/export/exportfile.go b/internal/export/exportfile.go
index fa1af85bc..02df64da5 100644
--- a/internal/export/exportfile.go
+++ b/internal/export/exportfile.go
@@ -67,6 +67,7 @@ func MarshalExportFile(eb *model.ExportBatch, exposures []*publishmodel.Exposure
 	// create compressed archive of binary and signature
 	buf := new(bytes.Buffer)
 	zw := zip.NewWriter(buf)
+
 	zf, err := zw.Create(exportBinaryName)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create zip entry for export: %w", err)
diff --git a/internal/util/generate.go b/internal/util/generate.go
index 270f39c45..d48f70257 100644
--- a/internal/util/generate.go
+++ b/internal/util/generate.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 	"github.com/google/exposure-notifications-server/testing/enclient"
+	"github.com/pkg/errors"
 )
 
 const (
@@ -109,6 +110,11 @@ func GenerateExposureKeys(numKeys, tr int, randomInterval bool) []v1alpha1.Expos
 	return exposureKeys
 }
 
+func CurrentIntervalNumber() int32 {
+	utcDay := time.Now().UTC().Truncate(24 * time.Hour)
+	return int32(utcDay.Unix() / 600)
+}
+
 // Creates a random exposure key.
 func RandomExposureKey(intervalNumber enclient.Interval, intervalCount int32, transmissionRisk int) (v1alpha1.ExposureKey, error) {
 	key, err := GenerateKey()
@@ -141,6 +147,15 @@ func GenerateKey() (string, error) {
 	return ToBase64(b), nil
 }
 
+func KeyFromString(source string) (string, error) {
+	if len(source) < dkLen {
+		return "", errors.Errorf("source '%v' should be at least %v len", source, dkLen)
+	}
+
+	b := []byte(source[:16])
+	return ToBase64(b), nil
+}
+
 // Encodes bytes array to base64.
 func ToBase64(key []byte) string {
 	return base64.StdEncoding.EncodeToString(key)
diff --git a/testing/enclient/enclient.go b/testing/enclient/enclient.go
index 63d61d211..50e022b6f 100644
--- a/testing/enclient/enclient.go
+++ b/testing/enclient/enclient.go
@@ -14,7 +14,7 @@ import (
 
 const (
 	// httpTimeout is the maximum amount of time to wait for a response.
-	httpTimeout = 30 * time.Second
+	httpTimeout = 5 * time.Minute
 )
 
 type Interval int32
diff --git a/testing/enclient/fs_client.go b/testing/enclient/fs_client.go
new file mode 100644
index 000000000..d6e8a4c7e
--- /dev/null
+++ b/testing/enclient/fs_client.go
@@ -0,0 +1,47 @@
+package enclient
+
+import (
+	"context"
+	"fmt"
+	"sort"
+
+	"cloud.google.com/go/storage"
+	"google.golang.org/api/iterator"
+)
+
+type CloudStorage struct {
+	client *storage.Client
+}
+
+func NewCloudStorage(ctx context.Context) (*CloudStorage, error) {
+	client, err := storage.NewClient(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("storage.NewClient: %w", err)
+	}
+	return &CloudStorage{client}, nil
+}
+
+func (gcs *CloudStorage) ListBucket(ctx context.Context, bucket string, filter func(attr *storage.ObjectAttrs) bool) []storage.ObjectAttrs {
+	query := &storage.Query{
+		Prefix: "exposureKeyExport-US/",
+	}
+	it := gcs.client.Bucket(bucket).Objects(ctx, query)
+
+	result := make([]storage.ObjectAttrs, 0)
+	for {
+		objAttrs, err := it.Next()
+		if err == iterator.Done {
+			break
+		}
+
+		if objAttrs != nil && filter(objAttrs) {
+			result = append(result, *objAttrs)
+		}
+	}
+
+	sort.SliceStable(result, func(i, j int) bool {
+		return result[j].Created.Before(result[i].Created)
+	})
+
+	return result
+}
diff --git a/testing/integration/en_api.go b/testing/integration/en_api.go
new file mode 100644
index 000000000..14f0ce653
--- /dev/null
+++ b/testing/integration/en_api.go
@@ -0,0 +1,134 @@
+// 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 integration
+
+import (
+	"context"
+	"io/ioutil"
+	"log"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+	"testing"
+
+	exportapi "github.com/google/exposure-notifications-server/internal/export"
+	"github.com/google/exposure-notifications-server/internal/logging"
+	"github.com/google/exposure-notifications-server/internal/monolith"
+	"github.com/google/exposure-notifications-server/internal/pb/export"
+	"github.com/google/exposure-notifications-server/internal/util"
+	"github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
+	"github.com/google/exposure-notifications-server/testing/enclient"
+)
+
+const appPackageName = "com.example.android.test"
+
+func collectExportResults(t *testing.T, exportDir string) *export.TemporaryExposureKeyExport {
+	exportFile := getExportFile(exportDir, t)
+
+	t.Logf("Reading keys data from: %v", exportFile)
+
+	blob, err := ioutil.ReadFile(exportFile)
+	if err != nil {
+		t.Fatalf("can't read export file: %v", err)
+	}
+
+	keyExport, err := exportapi.UnmarshalExportFile(blob)
+	if err != nil {
+		t.Fatalf("can't extract export data: %v", err)
+	}
+
+	return keyExport
+}
+
+func getExportFile(exportDir string, t *testing.T) string {
+	files, err := ioutil.ReadDir(exportDir)
+	if err != nil {
+		t.Fatalf("Can't read export directory: %v", err)
+	}
+
+	archiveFiles := make([]os.FileInfo, 0)
+	for _, f := range files {
+		if strings.HasSuffix(f.Name(), "zip") {
+			archiveFiles = append(archiveFiles, f)
+		}
+	}
+
+	if len(archiveFiles) < 1 {
+		t.Fatalf("can't find export archives in %v", exportDir)
+	}
+
+	sort.SliceStable(archiveFiles, func(i, j int) bool {
+		return files[i].Name() > files[j].Name()
+	})
+	exportFile := archiveFiles[0]
+	return filepath.Join(exportDir, exportFile.Name())
+}
+
+func publishKeys(t *testing.T, request v1alpha1.Publish) {
+	requestUrl := "http://localhost:8080/publish"
+
+	resp, err := enclient.PostRequest(requestUrl, request)
+	if err != nil {
+		t.Fatalf("request failed: %v, %v", err, resp)
+	}
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Publish request is sent to %v", requestUrl)
+}
+
+func exportBatches(t *testing.T) {
+	var bts []byte
+	requestUrl := "http://localhost:8080/export/create-batches"
+	resp, err := enclient.PostRequest(requestUrl, bts)
+	if err != nil {
+		t.Fatalf("request failed: %v, %v", err, resp)
+	}
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Create batches request is sent to %v", requestUrl)
+}
+
+func startExportWorkers(t *testing.T) {
+	var bts []byte
+	requestUrl := "http://localhost:8080/export/do-work"
+	resp, err := enclient.PostRequest(requestUrl, bts)
+	if err != nil {
+		t.Fatalf("request failed: %v, %v", err, resp)
+	}
+	log.Printf("response: %v", resp.Status)
+	t.Logf("Export worker request is sent to %v", requestUrl)
+}
+
+func publishRequest(keys []v1alpha1.ExposureKey, regions []string) v1alpha1.Publish {
+	padding, _ := util.RandomBytes(1000)
+	return v1alpha1.Publish{
+		Keys:                keys,
+		Regions:             regions,
+		AppPackageName:      appPackageName,
+		VerificationPayload: "Test Authority",
+		Padding:             util.ToBase64(padding),
+	}
+}
+
+func startServer() *monolith.MonoConfig {
+	ctx := context.Background()
+	logger := logging.FromContext(ctx)
+
+	config, err := monolith.RunServer(ctx)
+	if err != nil {
+		logger.Fatal(err)
+	}
+
+	return config
+}
diff --git a/testing/integration/publish_test.go b/testing/integration/publish_test.go
index ede48f6d6..9d619bdc9 100644
--- a/testing/integration/publish_test.go
+++ b/testing/integration/publish_test.go
@@ -1,62 +1,89 @@
 package integration
 
 import (
-	"log"
+	"encoding/json"
 	"testing"
+	"time"
 
-	"github.com/google/exposure-notifications-server/internal/database"
-	"github.com/google/exposure-notifications-server/testing/enclient"
+	"github.com/golang/protobuf/proto"
+	export2 "github.com/google/exposure-notifications-server/internal/pb/export"
+	"github.com/google/exposure-notifications-server/internal/util"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
-const appPackageName = ""
+const (
+	keysInRequest    = 14
+	transmissionRisk = 5
+	region           = "TEST"
+	exportDir        = "/tmp/en/exposureKeyExport-e2e"
+)
 
-func TestCleanupAPI(t *testing.T) {
-	var bts []byte
-	requestUrl := "http://localhost:8080/cleanup-exposure"
-	resp, err := enclient.PostRequest(requestUrl, bts)
-	if err != nil {
-		t.Errorf("request failed: %v, %v", err, resp)
-		return
-	}
+func TestPublish(t *testing.T) {
+	keys := util.GenerateExposureKeys(keysInRequest, transmissionRisk, true)
+	request := publishRequest(keys, []string{region})
 
-	log.Printf("response: %v", resp.Status)
-	t.Logf("Cleanup request is sent to %v", requestUrl)
-}
+	publishKeys(t, request)
 
-func TestExportAPI(t *testing.T) {
-	var bts []byte
-	requestUrl := "http://localhost:8080/export/create-batches"
-	resp, err := enclient.PostRequest(requestUrl, bts)
-	if err != nil {
-		t.Errorf("request failed: %v, %v", err, resp)
-		return
+	t.Logf("Waiting before creating batches.")
+	time.Sleep(1 * time.Minute)
+
+	exportBatches(t)
+
+	t.Logf("Waiting before starting export workers.")
+	time.Sleep(5 * time.Second)
+
+	startExportWorkers(t)
+
+	got := collectExportResults(t, exportDir)
+
+	wantedKeysMap := make(map[string]export2.TemporaryExposureKey)
+	for _, key := range keys {
+		wantedKeysMap[key.Key] = export2.TemporaryExposureKey{
+			KeyData:                    util.DecodeKey(key.Key),
+			TransmissionRiskLevel:      proto.Int32(int32(key.TransmissionRisk)),
+			RollingStartIntervalNumber: proto.Int32(key.IntervalNumber),
+			RollingPeriod:              proto.Int32(key.IntervalCount),
+		}
 	}
 
-	log.Printf("response: %v", resp.Status)
-	t.Logf("Create batches request is sent to %v", requestUrl)
-}
+	want := export2.TemporaryExposureKeyExport{
+		StartTimestamp: nil,
+		EndTimestamp:   nil,
+		Region:         proto.String("TEST"),
+		BatchNum:       proto.Int32(1),
+		BatchSize:      proto.Int32(1),
+		SignatureInfos: nil,
+		Keys:           nil,
+	}
 
-func TestExportWorkerApi(t *testing.T) {
-	var bts []byte
-	requestUrl := "http://localhost:8080/export/do-work"
-	resp, err := enclient.PostRequest(requestUrl, bts)
-	if err != nil {
-		t.Errorf("request failed: %v, %v", err, resp)
-		return
+	options := []cmp.Option{
+		cmpopts.IgnoreFields(want, "StartTimestamp"),
+		cmpopts.IgnoreFields(want, "EndTimestamp"),
+		cmpopts.IgnoreFields(want, "SignatureInfos"),
+		cmpopts.IgnoreFields(want, "Keys"),
+		cmpopts.IgnoreUnexported(want),
 	}
 
-	log.Printf("response: %v", resp.Status)
-	t.Logf("Export worker request is sent to %v", requestUrl)
-}
+	diff := cmp.Diff(got, &want, options...)
+	if diff != "" {
+		t.Errorf("%v", diff)
+	}
 
-func publishKey(keys []database.ExposureKey, regions []string) database.Publish {
-	padding := enclient.RandomBytes(1000)
-	return database.Publish{
-		Keys:                      keys,
-		Regions:                   regions,
-		AppPackageName:            appPackageName,
-		DeviceVerificationPayload: "Test Device Verification Payload",
-		VerificationPayload:       "Test Authority",
-		Padding:                   enclient.ToBase64(padding),
+	for _, key := range got.Keys {
+		s := util.ToBase64(key.KeyData)
+		wantedKey := wantedKeysMap[s]
+		gotKey := *key
+		diff := cmp.Diff(wantedKey, gotKey, cmpopts.IgnoreUnexported(gotKey))
+		if diff != "" {
+			t.Errorf("invalid key value: %v:%v", s, diff)
+		}
 	}
+
+	bytes, err := json.MarshalIndent(got, "", "")
+	if err != nil {
+		t.Fatalf("can't marshal json results: %v", err)
+	}
+
+	t.Logf("%v", string(bytes))
 }

From 221ca18d8bec27380d573a6712237ba407eabb9e Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Tue, 9 Jun 2020 18:58:52 -0700
Subject: [PATCH 03/12] Test Code Cleanup

---
 go.mod                        |  1 -
 internal/export/exportfile.go |  1 -
 internal/util/generate.go     | 15 -----------
 testing/enclient/fs_client.go | 47 -----------------------------------
 4 files changed, 64 deletions(-)
 delete mode 100644 testing/enclient/fs_client.go

diff --git a/go.mod b/go.mod
index 2ed222bff..afe42fdd7 100644
--- a/go.mod
+++ b/go.mod
@@ -54,7 +54,6 @@ require (
 	github.com/oracle/oci-go-sdk v19.3.0+incompatible // indirect
 	github.com/ory/dockertest v3.3.5+incompatible
 	github.com/pierrec/lz4 v2.5.2+incompatible // indirect
-	github.com/pkg/errors v0.9.1
 	github.com/posener/complete v1.2.3 // indirect
 	github.com/prometheus/client_golang v1.6.0 // indirect
 	github.com/prometheus/common v0.10.0 // indirect
diff --git a/internal/export/exportfile.go b/internal/export/exportfile.go
index 02df64da5..fa1af85bc 100644
--- a/internal/export/exportfile.go
+++ b/internal/export/exportfile.go
@@ -67,7 +67,6 @@ func MarshalExportFile(eb *model.ExportBatch, exposures []*publishmodel.Exposure
 	// create compressed archive of binary and signature
 	buf := new(bytes.Buffer)
 	zw := zip.NewWriter(buf)
-
 	zf, err := zw.Create(exportBinaryName)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create zip entry for export: %w", err)
diff --git a/internal/util/generate.go b/internal/util/generate.go
index d48f70257..270f39c45 100644
--- a/internal/util/generate.go
+++ b/internal/util/generate.go
@@ -24,7 +24,6 @@ import (
 
 	"github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 	"github.com/google/exposure-notifications-server/testing/enclient"
-	"github.com/pkg/errors"
 )
 
 const (
@@ -110,11 +109,6 @@ func GenerateExposureKeys(numKeys, tr int, randomInterval bool) []v1alpha1.Expos
 	return exposureKeys
 }
 
-func CurrentIntervalNumber() int32 {
-	utcDay := time.Now().UTC().Truncate(24 * time.Hour)
-	return int32(utcDay.Unix() / 600)
-}
-
 // Creates a random exposure key.
 func RandomExposureKey(intervalNumber enclient.Interval, intervalCount int32, transmissionRisk int) (v1alpha1.ExposureKey, error) {
 	key, err := GenerateKey()
@@ -147,15 +141,6 @@ func GenerateKey() (string, error) {
 	return ToBase64(b), nil
 }
 
-func KeyFromString(source string) (string, error) {
-	if len(source) < dkLen {
-		return "", errors.Errorf("source '%v' should be at least %v len", source, dkLen)
-	}
-
-	b := []byte(source[:16])
-	return ToBase64(b), nil
-}
-
 // Encodes bytes array to base64.
 func ToBase64(key []byte) string {
 	return base64.StdEncoding.EncodeToString(key)
diff --git a/testing/enclient/fs_client.go b/testing/enclient/fs_client.go
deleted file mode 100644
index d6e8a4c7e..000000000
--- a/testing/enclient/fs_client.go
+++ /dev/null
@@ -1,47 +0,0 @@
-package enclient
-
-import (
-	"context"
-	"fmt"
-	"sort"
-
-	"cloud.google.com/go/storage"
-	"google.golang.org/api/iterator"
-)
-
-type CloudStorage struct {
-	client *storage.Client
-}
-
-func NewCloudStorage(ctx context.Context) (*CloudStorage, error) {
-	client, err := storage.NewClient(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("storage.NewClient: %w", err)
-	}
-	return &CloudStorage{client}, nil
-}
-
-func (gcs *CloudStorage) ListBucket(ctx context.Context, bucket string, filter func(attr *storage.ObjectAttrs) bool) []storage.ObjectAttrs {
-	query := &storage.Query{
-		Prefix: "exposureKeyExport-US/",
-	}
-	it := gcs.client.Bucket(bucket).Objects(ctx, query)
-
-	result := make([]storage.ObjectAttrs, 0)
-	for {
-		objAttrs, err := it.Next()
-		if err == iterator.Done {
-			break
-		}
-
-		if objAttrs != nil && filter(objAttrs) {
-			result = append(result, *objAttrs)
-		}
-	}
-
-	sort.SliceStable(result, func(i, j int) bool {
-		return result[j].Created.Before(result[i].Created)
-	})
-
-	return result
-}

From cc460e55ca3d89643d50b819c5788dcd8875ace0 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Tue, 9 Jun 2020 23:13:32 -0700
Subject: [PATCH 04/12] Integrating test code with the in-process SUT.

---
 {testing => internal}/integration/en_api.go |  86 ++++++----
 internal/integration/integration_test.go    |  16 +-
 internal/integration/publish_test.go        | 169 ++++++++++++++++----
 internal/storage/memory.go                  |  16 ++
 testing/integration/publish_test.go         |  89 -----------
 5 files changed, 208 insertions(+), 168 deletions(-)
 rename {testing => internal}/integration/en_api.go (60%)
 delete mode 100644 testing/integration/publish_test.go

diff --git a/testing/integration/en_api.go b/internal/integration/en_api.go
similarity index 60%
rename from testing/integration/en_api.go
rename to internal/integration/en_api.go
index 14f0ce653..eaf2e9654 100644
--- a/testing/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -15,9 +15,12 @@
 package integration
 
 import (
-	"context"
+	"bytes"
+	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"os"
 	"path/filepath"
 	"sort"
@@ -25,12 +28,9 @@ import (
 	"testing"
 
 	exportapi "github.com/google/exposure-notifications-server/internal/export"
-	"github.com/google/exposure-notifications-server/internal/logging"
-	"github.com/google/exposure-notifications-server/internal/monolith"
 	"github.com/google/exposure-notifications-server/internal/pb/export"
 	"github.com/google/exposure-notifications-server/internal/util"
-	"github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
-	"github.com/google/exposure-notifications-server/testing/enclient"
+	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 )
 
 const appPackageName = "com.example.android.test"
@@ -53,6 +53,46 @@ func collectExportResults(t *testing.T, exportDir string) *export.TemporaryExpos
 	return keyExport
 }
 
+type EnServerClient struct {
+	client *http.Client
+}
+
+// Posts requests to the specified url.
+// This methods attempts to serialize data argument as a json.
+func (server EnServerClient) postRequest(url string, data interface{}) (*http.Response, error) {
+	request := bytes.NewBuffer(JSONRequest(data))
+	r, err := http.NewRequest("POST", url, request)
+	if err != nil {
+		return nil, err
+	}
+	r.Header.Set("Content-Type", "application/json")
+	resp, err := server.client.Do(r)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		// Return error upstream.
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to copy error body (%d): %w", resp.StatusCode, err)
+		}
+		return resp, fmt.Errorf("post request failed with status: %v\n%v", resp.StatusCode, body)
+	}
+
+	return resp, nil
+}
+
+// Serializes the given argument to json.
+func JSONRequest(data interface{}) []byte {
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		log.Fatalf("unable to marshal json payload")
+	}
+	return jsonData
+}
+
 func getExportFile(exportDir string, t *testing.T) string {
 	files, err := ioutil.ReadDir(exportDir)
 	if err != nil {
@@ -77,21 +117,19 @@ func getExportFile(exportDir string, t *testing.T) string {
 	return filepath.Join(exportDir, exportFile.Name())
 }
 
-func publishKeys(t *testing.T, request v1alpha1.Publish) {
-	requestUrl := "http://localhost:8080/publish"
-
-	resp, err := enclient.PostRequest(requestUrl, request)
+func (enServer EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish) {
+	resp, err := enServer.postRequest("/publish", request)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
 	log.Printf("response: %v", resp.Status)
-	t.Logf("Publish request is sent to %v", requestUrl)
+	t.Logf("Publish request is sent to %v", "/publish")
 }
 
-func exportBatches(t *testing.T) {
+func (enServer EnServerClient) ExportBatches(t *testing.T) {
 	var bts []byte
-	requestUrl := "http://localhost:8080/export/create-batches"
-	resp, err := enclient.PostRequest(requestUrl, bts)
+	requestUrl := "/export/create-batches"
+	resp, err := enServer.postRequest(requestUrl, bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
@@ -99,10 +137,10 @@ func exportBatches(t *testing.T) {
 	t.Logf("Create batches request is sent to %v", requestUrl)
 }
 
-func startExportWorkers(t *testing.T) {
+func (enServer EnServerClient) StartExportWorkers(t *testing.T) {
 	var bts []byte
-	requestUrl := "http://localhost:8080/export/do-work"
-	resp, err := enclient.PostRequest(requestUrl, bts)
+	requestUrl := "/export/do-work"
+	resp, err := enServer.postRequest(requestUrl, bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
@@ -110,9 +148,9 @@ func startExportWorkers(t *testing.T) {
 	t.Logf("Export worker request is sent to %v", requestUrl)
 }
 
-func publishRequest(keys []v1alpha1.ExposureKey, regions []string) v1alpha1.Publish {
+func publishRequest(keys []verifyapi.ExposureKey, regions []string) verifyapi.Publish {
 	padding, _ := util.RandomBytes(1000)
-	return v1alpha1.Publish{
+	return verifyapi.Publish{
 		Keys:                keys,
 		Regions:             regions,
 		AppPackageName:      appPackageName,
@@ -120,15 +158,3 @@ func publishRequest(keys []v1alpha1.ExposureKey, regions []string) v1alpha1.Publ
 		Padding:             util.ToBase64(padding),
 	}
 }
-
-func startServer() *monolith.MonoConfig {
-	ctx := context.Background()
-	logger := logging.FromContext(ctx)
-
-	config, err := monolith.RunServer(ctx)
-	if err != nil {
-		logger.Fatal(err)
-	}
-
-	return config
-}
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index d54997fb7..e52833d63 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -33,7 +33,7 @@ import (
 	"github.com/google/exposure-notifications-server/internal/storage"
 )
 
-func testServer(tb testing.TB) (*serverenv.ServerEnv, *http.Client) {
+func testServer(tb testing.TB, exportConfig *export.Config) (*serverenv.ServerEnv, *http.Client) {
 	tb.Helper()
 
 	ctx := context.Background()
@@ -96,18 +96,6 @@ func testServer(tb testing.TB) (*serverenv.ServerEnv, *http.Client) {
 	}
 	mux.Handle("/cleanup-exposure", cleanupExposureHandler)
 
-	// Export
-	exportConfig := &export.Config{
-		CreateTimeout:  10 * time.Second,
-		WorkerTimeout:  10 * time.Second,
-		MinRecords:     1,
-		PaddingRange:   1,
-		MaxRecords:     10000,
-		TruncateWindow: 1 * time.Second,
-		MinWindowAge:   1 * time.Second,
-		TTL:            336 * time.Hour,
-	}
-
 	exportServer, err := export.NewServer(exportConfig, env)
 	if err != nil {
 		tb.Fatal(err)
@@ -131,7 +119,7 @@ func testServer(tb testing.TB) (*serverenv.ServerEnv, *http.Client) {
 		MinRequestDuration:      50 * time.Millisecond,
 		MaxKeysOnPublish:        15,
 		MaxIntervalAge:          360 * time.Hour,
-		TruncateWindow:          1 * time.Hour,
+		TruncateWindow:          1 * time.Second,
 		DebugAPIResponses:       true,
 		DebugReleaseSameDayKeys: true,
 	}
diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 147dcd392..826c5a3d7 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -15,35 +15,56 @@
 package integration
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
-	"net/http"
+	"sort"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/golang/protobuf/proto"
 	authorizedappmodel "github.com/google/exposure-notifications-server/internal/authorizedapp/model"
+	exportapi "github.com/google/exposure-notifications-server/internal/export"
 	exportdatabase "github.com/google/exposure-notifications-server/internal/export/database"
 	exportmodel "github.com/google/exposure-notifications-server/internal/export/model"
+	"github.com/google/exposure-notifications-server/internal/pb/export"
 	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/storage"
 	"github.com/google/exposure-notifications-server/internal/util"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 )
 
 func TestPublish(t *testing.T) {
 	t.Parallel()
 
 	ctx := context.Background()
-	env, client := testServer(t)
+
+	// Export
+	exportConfig := &exportapi.Config{
+		CreateTimeout:  10 * time.Second,
+		WorkerTimeout:  10 * time.Second,
+		MinRecords:     1,
+		PaddingRange:   1,
+		MaxRecords:     10000,
+		TruncateWindow: 1 * time.Millisecond,
+		MinWindowAge:   1 * time.Second,
+		TTL:            336 * time.Hour,
+	}
+
+	env, client := testServer(t, exportConfig)
 	db := env.Database()
+	server := &EnServerClient{client: client}
 
 	// Create an authorized app.
 	aa := env.AuthorizedAppProvider()
 	if err := aa.Add(ctx, &authorizedappmodel.AuthorizedApp{
 		AppPackageName: "com.example.app",
 		AllowedRegions: map[string]struct{}{
-			"US": {},
+			"TEST": {},
 		},
 		AllowedHealthAuthorityIDs: map[int64]struct{}{
 			12345: {},
@@ -66,42 +87,29 @@ func TestPublish(t *testing.T) {
 	}
 
 	// Create an export config.
+	exportPeriod := 2 * time.Second
 	ec := &exportmodel.ExportConfig{
 		BucketName:       "my-bucket",
-		Period:           1 * time.Second,
-		OutputRegion:     "US",
+		Period:           exportPeriod,
+		OutputRegion:     "TEST",
 		From:             time.Now().Add(-2 * time.Second),
 		Thru:             time.Now().Add(1 * time.Hour),
-		SignatureInfoIDs: []int64{si.ID},
+		SignatureInfoIDs: []int64{},
 	}
 	if err := exportdatabase.New(db).AddExportConfig(ctx, ec); err != nil {
 		t.Fatal(err)
 	}
 
-	payload := &verifyapi.Publish{
+	payload := verifyapi.Publish{
 		Keys:           util.GenerateExposureKeys(3, -1, false),
-		Regions:        []string{"US"},
+		Regions:        []string{"TEST"},
 		AppPackageName: "com.example.app",
 
 		// TODO: hook up verification
 		VerificationPayload: "TODO",
 	}
 
-	var body bytes.Buffer
-	if err := json.NewEncoder(&body).Encode(payload); err != nil {
-		t.Fatal(err)
-	}
-
-	resp, err := client.Post("/publish", "application/json", &body)
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer resp.Body.Close()
-
-	// Ensure we get a successful response code.
-	if got, want := resp.StatusCode, http.StatusOK; got != want {
-		t.Errorf("expected %v to be %v", got, want)
-	}
+	server.PublishKeys(t, payload)
 
 	// Look up the exposures in the database.
 	criteria := publishdb.IterateExposuresCriteria{
@@ -110,6 +118,7 @@ func TestPublish(t *testing.T) {
 
 	var exposures []*publishmodel.Exposure
 	if _, err := publishdb.New(db).IterateExposures(ctx, criteria, func(m *publishmodel.Exposure) error {
+		t.Logf("NEW EXPOSURE: %v", m)
 		exposures = append(exposures, m)
 		return nil
 	}); err != nil {
@@ -120,25 +129,115 @@ func TestPublish(t *testing.T) {
 		t.Errorf("expected %v to be %v: %#v", got, want, exposures)
 	}
 
-	// Create an export.
-	resp, err = client.Get("/export/create-batches")
+	wait(t, exportPeriod+500*time.Millisecond, "Waiting before export batches")
+	server.ExportBatches(t)
+
+	wait(t, 500*time.Millisecond, "Waiting before staring workers")
+	server.StartExportWorkers(t)
+
+	memory, ok := env.Blobstore().(*storage.Memory)
+	if !ok {
+		t.Fatalf("can't use %t blobstore for verification", env.Blobstore())
+	}
+	keyExport := getKeysFromLatestBatch(t, "my-bucket", ctx, env, memory)
+
+	got := keyExport
+
+	wantedKeysMap := make(map[string]export.TemporaryExposureKey)
+	for _, key := range payload.Keys {
+		wantedKeysMap[key.Key] = export.TemporaryExposureKey{
+			KeyData:                    util.DecodeKey(key.Key),
+			TransmissionRiskLevel:      proto.Int32(int32(key.TransmissionRisk)),
+			RollingStartIntervalNumber: proto.Int32(key.IntervalNumber),
+		}
+	}
+
+	want := export.TemporaryExposureKeyExport{
+		StartTimestamp: nil,
+		EndTimestamp:   nil,
+		Region:         proto.String("TEST"),
+		BatchNum:       proto.Int32(1),
+		BatchSize:      proto.Int32(1),
+		SignatureInfos: nil,
+		Keys:           nil,
+	}
+
+	options := []cmp.Option{
+		cmpopts.IgnoreFields(want, "StartTimestamp"),
+		cmpopts.IgnoreFields(want, "EndTimestamp"),
+		cmpopts.IgnoreFields(want, "SignatureInfos"),
+		cmpopts.IgnoreFields(want, "Keys"),
+		cmpopts.IgnoreUnexported(want),
+	}
+
+	diff := cmp.Diff(got, &want, options...)
+	if diff != "" {
+		t.Errorf("%v", diff)
+	}
+
+	for _, key := range got.Keys {
+		s := util.ToBase64(key.KeyData)
+		wantedKey := wantedKeysMap[s]
+		gotKey := *key
+		diff := cmp.Diff(wantedKey, gotKey, cmpopts.IgnoreUnexported(gotKey))
+		if diff != "" {
+
+			t.Logf("WANT: %v", proto.MarshalTextString(&wantedKey))
+			t.Logf(" GOT: %v", proto.MarshalTextString(&gotKey))
+
+			t.Errorf("invalid key value: %v:%v", s, diff)
+		}
+	}
+
+	bytes, err := json.MarshalIndent(got, "", "")
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("can't marshal json results: %v", err)
 	}
-	defer resp.Body.Close()
 
-	resp, err = client.Get("/export/do-work")
+	t.Logf("%v", string(bytes))
+	// TODO: verify signature
+}
+
+func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context, env *serverenv.ServerEnv, memory *storage.Memory) *export.TemporaryExposureKeyExport {
+	exportFile := getLatestFile(t, memory, ctx, exportDir)
+
+	t.Logf("Reading keys data from: %v", exportFile)
+
+	blob, err := env.Blobstore().GetObject(ctx, "", exportFile)
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer resp.Body.Close()
 
-	// TODO: verify export has the correct file
-	b, err := env.Blobstore().GetObject(ctx, "my-bucket", "index.txt")
+	keyExport, err := exportapi.UnmarshalExportFile(blob)
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("can't extract export data: %v", err)
 	}
-	_ = b
 
-	// TODO: verify signature
+	return keyExport
+}
+
+func getLatestFile(t *testing.T, blobstore *storage.Memory, ctx context.Context, exportDir string) string {
+	files := blobstore.ListObjects(ctx, exportDir)
+
+	archiveFiles := make([]string, 0)
+	for fileName := range files {
+		if strings.HasSuffix(fileName, "zip") {
+			archiveFiles = append(archiveFiles, fileName)
+		}
+	}
+
+	if len(archiveFiles) < 1 {
+		t.Fatalf("can't find export archives in %v", exportDir)
+	}
+
+	sort.SliceStable(archiveFiles, func(i, j int) bool {
+		return archiveFiles[i] > archiveFiles[j]
+	})
+	exportFile := archiveFiles[0]
+	return exportFile
+}
+
+func wait(t *testing.T, duration time.Duration, message string) {
+	t.Logf("%s - waiting for %v", message, duration)
+	time.Sleep(duration)
 }
diff --git a/internal/storage/memory.go b/internal/storage/memory.go
index 83ea547a3..a1a7805ac 100644
--- a/internal/storage/memory.go
+++ b/internal/storage/memory.go
@@ -17,6 +17,7 @@ package storage
 import (
 	"context"
 	"path"
+	"strings"
 	"sync"
 )
 
@@ -71,3 +72,18 @@ func (s *Memory) GetObject(_ context.Context, folder, filename string) ([]byte,
 	}
 	return v, nil
 }
+
+// GetObject returns the contents for the given object. If the object does not
+// exist, it returns ErrNotFound.
+func (s *Memory) ListObjects(_ context.Context, folder string) map[string][]byte {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	result := make(map[string][]byte)
+	for k, v := range s.data {
+		if strings.HasPrefix(k, folder+"/") {
+			result[k] = v
+		}
+	}
+	return result
+}
diff --git a/testing/integration/publish_test.go b/testing/integration/publish_test.go
deleted file mode 100644
index 9d619bdc9..000000000
--- a/testing/integration/publish_test.go
+++ /dev/null
@@ -1,89 +0,0 @@
-package integration
-
-import (
-	"encoding/json"
-	"testing"
-	"time"
-
-	"github.com/golang/protobuf/proto"
-	export2 "github.com/google/exposure-notifications-server/internal/pb/export"
-	"github.com/google/exposure-notifications-server/internal/util"
-	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-cmp/cmp/cmpopts"
-)
-
-const (
-	keysInRequest    = 14
-	transmissionRisk = 5
-	region           = "TEST"
-	exportDir        = "/tmp/en/exposureKeyExport-e2e"
-)
-
-func TestPublish(t *testing.T) {
-	keys := util.GenerateExposureKeys(keysInRequest, transmissionRisk, true)
-	request := publishRequest(keys, []string{region})
-
-	publishKeys(t, request)
-
-	t.Logf("Waiting before creating batches.")
-	time.Sleep(1 * time.Minute)
-
-	exportBatches(t)
-
-	t.Logf("Waiting before starting export workers.")
-	time.Sleep(5 * time.Second)
-
-	startExportWorkers(t)
-
-	got := collectExportResults(t, exportDir)
-
-	wantedKeysMap := make(map[string]export2.TemporaryExposureKey)
-	for _, key := range keys {
-		wantedKeysMap[key.Key] = export2.TemporaryExposureKey{
-			KeyData:                    util.DecodeKey(key.Key),
-			TransmissionRiskLevel:      proto.Int32(int32(key.TransmissionRisk)),
-			RollingStartIntervalNumber: proto.Int32(key.IntervalNumber),
-			RollingPeriod:              proto.Int32(key.IntervalCount),
-		}
-	}
-
-	want := export2.TemporaryExposureKeyExport{
-		StartTimestamp: nil,
-		EndTimestamp:   nil,
-		Region:         proto.String("TEST"),
-		BatchNum:       proto.Int32(1),
-		BatchSize:      proto.Int32(1),
-		SignatureInfos: nil,
-		Keys:           nil,
-	}
-
-	options := []cmp.Option{
-		cmpopts.IgnoreFields(want, "StartTimestamp"),
-		cmpopts.IgnoreFields(want, "EndTimestamp"),
-		cmpopts.IgnoreFields(want, "SignatureInfos"),
-		cmpopts.IgnoreFields(want, "Keys"),
-		cmpopts.IgnoreUnexported(want),
-	}
-
-	diff := cmp.Diff(got, &want, options...)
-	if diff != "" {
-		t.Errorf("%v", diff)
-	}
-
-	for _, key := range got.Keys {
-		s := util.ToBase64(key.KeyData)
-		wantedKey := wantedKeysMap[s]
-		gotKey := *key
-		diff := cmp.Diff(wantedKey, gotKey, cmpopts.IgnoreUnexported(gotKey))
-		if diff != "" {
-			t.Errorf("invalid key value: %v:%v", s, diff)
-		}
-	}
-
-	bytes, err := json.MarshalIndent(got, "", "")
-	if err != nil {
-		t.Fatalf("can't marshal json results: %v", err)
-	}
-
-	t.Logf("%v", string(bytes))
-}

From 8c311f17437271b5d8051f609519171c374753e1 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 10 Jun 2020 17:21:33 -0700
Subject: [PATCH 05/12] Addressing PR review comments.

---
 internal/integration/en_api.go       | 89 ++++------------------------
 internal/integration/publish_test.go | 48 +++++++--------
 internal/storage/memory.go           |  3 +-
 testing/enclient/enclient.go         |  2 +-
 4 files changed, 33 insertions(+), 109 deletions(-)

diff --git a/internal/integration/en_api.go b/internal/integration/en_api.go
index eaf2e9654..fb2700599 100644
--- a/internal/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -21,38 +21,11 @@ import (
 	"io/ioutil"
 	"log"
 	"net/http"
-	"os"
-	"path/filepath"
-	"sort"
-	"strings"
 	"testing"
 
-	exportapi "github.com/google/exposure-notifications-server/internal/export"
-	"github.com/google/exposure-notifications-server/internal/pb/export"
-	"github.com/google/exposure-notifications-server/internal/util"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 )
 
-const appPackageName = "com.example.android.test"
-
-func collectExportResults(t *testing.T, exportDir string) *export.TemporaryExposureKeyExport {
-	exportFile := getExportFile(exportDir, t)
-
-	t.Logf("Reading keys data from: %v", exportFile)
-
-	blob, err := ioutil.ReadFile(exportFile)
-	if err != nil {
-		t.Fatalf("can't read export file: %v", err)
-	}
-
-	keyExport, err := exportapi.UnmarshalExportFile(blob)
-	if err != nil {
-		t.Fatalf("can't extract export data: %v", err)
-	}
-
-	return keyExport
-}
-
 type EnServerClient struct {
 	client *http.Client
 }
@@ -60,7 +33,11 @@ type EnServerClient struct {
 // Posts requests to the specified url.
 // This methods attempts to serialize data argument as a json.
 func (server EnServerClient) postRequest(url string, data interface{}) (*http.Response, error) {
-	request := bytes.NewBuffer(JSONRequest(data))
+	jsonData, err := json.Marshal(data)
+	if err != nil {
+		return nil, fmt.Errorf("unable to marshal json payload")
+	}
+	request := bytes.NewBuffer(jsonData)
 	r, err := http.NewRequest("POST", url, request)
 	if err != nil {
 		return nil, err
@@ -84,41 +61,8 @@ func (server EnServerClient) postRequest(url string, data interface{}) (*http.Re
 	return resp, nil
 }
 
-// Serializes the given argument to json.
-func JSONRequest(data interface{}) []byte {
-	jsonData, err := json.Marshal(data)
-	if err != nil {
-		log.Fatalf("unable to marshal json payload")
-	}
-	return jsonData
-}
-
-func getExportFile(exportDir string, t *testing.T) string {
-	files, err := ioutil.ReadDir(exportDir)
-	if err != nil {
-		t.Fatalf("Can't read export directory: %v", err)
-	}
-
-	archiveFiles := make([]os.FileInfo, 0)
-	for _, f := range files {
-		if strings.HasSuffix(f.Name(), "zip") {
-			archiveFiles = append(archiveFiles, f)
-		}
-	}
-
-	if len(archiveFiles) < 1 {
-		t.Fatalf("can't find export archives in %v", exportDir)
-	}
-
-	sort.SliceStable(archiveFiles, func(i, j int) bool {
-		return files[i].Name() > files[j].Name()
-	})
-	exportFile := archiveFiles[0]
-	return filepath.Join(exportDir, exportFile.Name())
-}
-
-func (enServer EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish) {
-	resp, err := enServer.postRequest("/publish", request)
+func (server EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish) {
+	resp, err := server.postRequest("/publish", request)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
@@ -126,10 +70,10 @@ func (enServer EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publi
 	t.Logf("Publish request is sent to %v", "/publish")
 }
 
-func (enServer EnServerClient) ExportBatches(t *testing.T) {
+func (server EnServerClient) ExportBatches(t *testing.T) {
 	var bts []byte
 	requestUrl := "/export/create-batches"
-	resp, err := enServer.postRequest(requestUrl, bts)
+	resp, err := server.postRequest(requestUrl, bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
@@ -137,24 +81,13 @@ func (enServer EnServerClient) ExportBatches(t *testing.T) {
 	t.Logf("Create batches request is sent to %v", requestUrl)
 }
 
-func (enServer EnServerClient) StartExportWorkers(t *testing.T) {
+func (server EnServerClient) StartExportWorkers(t *testing.T) {
 	var bts []byte
 	requestUrl := "/export/do-work"
-	resp, err := enServer.postRequest(requestUrl, bts)
+	resp, err := server.postRequest(requestUrl, bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
 	log.Printf("response: %v", resp.Status)
 	t.Logf("Export worker request is sent to %v", requestUrl)
 }
-
-func publishRequest(keys []verifyapi.ExposureKey, regions []string) verifyapi.Publish {
-	padding, _ := util.RandomBytes(1000)
-	return verifyapi.Publish{
-		Keys:                keys,
-		Regions:             regions,
-		AppPackageName:      appPackageName,
-		VerificationPayload: "Test Authority",
-		Padding:             util.ToBase64(padding),
-	}
-}
diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 826c5a3d7..d7d47b699 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -17,7 +17,6 @@ package integration
 import (
 	"context"
 	"encoding/json"
-	"sort"
 	"strings"
 	"testing"
 	"time"
@@ -57,7 +56,7 @@ func TestPublish(t *testing.T) {
 
 	env, client := testServer(t, exportConfig)
 	db := env.Database()
-	server := &EnServerClient{client: client}
+	enClient := &EnServerClient{client: client}
 
 	// Create an authorized app.
 	aa := env.AuthorizedAppProvider()
@@ -109,7 +108,7 @@ func TestPublish(t *testing.T) {
 		VerificationPayload: "TODO",
 	}
 
-	server.PublishKeys(t, payload)
+	enClient.PublishKeys(t, payload)
 
 	// Look up the exposures in the database.
 	criteria := publishdb.IterateExposuresCriteria{
@@ -118,7 +117,6 @@ func TestPublish(t *testing.T) {
 
 	var exposures []*publishmodel.Exposure
 	if _, err := publishdb.New(db).IterateExposures(ctx, criteria, func(m *publishmodel.Exposure) error {
-		t.Logf("NEW EXPOSURE: %v", m)
 		exposures = append(exposures, m)
 		return nil
 	}); err != nil {
@@ -129,11 +127,13 @@ func TestPublish(t *testing.T) {
 		t.Errorf("expected %v to be %v: %#v", got, want, exposures)
 	}
 
-	wait(t, exportPeriod+500*time.Millisecond, "Waiting before export batches")
-	server.ExportBatches(t)
+	t.Logf("Waiting %v before export batches", exportPeriod+500*time.Millisecond)
+	time.Sleep(exportPeriod + 500*time.Millisecond)
+	enClient.ExportBatches(t)
 
-	wait(t, 500*time.Millisecond, "Waiting before staring workers")
-	server.StartExportWorkers(t)
+	t.Logf("Waiting %v before starting workers", 500*time.Millisecond)
+	time.Sleep(500 * time.Millisecond)
+	enClient.StartExportWorkers(t)
 
 	memory, ok := env.Blobstore().(*storage.Memory)
 	if !ok {
@@ -181,10 +181,6 @@ func TestPublish(t *testing.T) {
 		gotKey := *key
 		diff := cmp.Diff(wantedKey, gotKey, cmpopts.IgnoreUnexported(gotKey))
 		if diff != "" {
-
-			t.Logf("WANT: %v", proto.MarshalTextString(&wantedKey))
-			t.Logf(" GOT: %v", proto.MarshalTextString(&gotKey))
-
 			t.Errorf("invalid key value: %v:%v", s, diff)
 		}
 	}
@@ -200,6 +196,9 @@ func TestPublish(t *testing.T) {
 
 func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context, env *serverenv.ServerEnv, memory *storage.Memory) *export.TemporaryExposureKeyExport {
 	exportFile := getLatestFile(t, memory, ctx, exportDir)
+	if exportFile == "" {
+		t.Fatalf("Can't find export files in blobstore: %v", exportDir)
+	}
 
 	t.Logf("Reading keys data from: %v", exportFile)
 
@@ -219,25 +218,18 @@ func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context,
 func getLatestFile(t *testing.T, blobstore *storage.Memory, ctx context.Context, exportDir string) string {
 	files := blobstore.ListObjects(ctx, exportDir)
 
-	archiveFiles := make([]string, 0)
+	latestFileName := ""
 	for fileName := range files {
 		if strings.HasSuffix(fileName, "zip") {
-			archiveFiles = append(archiveFiles, fileName)
+			if latestFileName == "" {
+				latestFileName = fileName
+			} else {
+				if fileName > latestFileName {
+					latestFileName = fileName
+				}
+			}
 		}
 	}
 
-	if len(archiveFiles) < 1 {
-		t.Fatalf("can't find export archives in %v", exportDir)
-	}
-
-	sort.SliceStable(archiveFiles, func(i, j int) bool {
-		return archiveFiles[i] > archiveFiles[j]
-	})
-	exportFile := archiveFiles[0]
-	return exportFile
-}
-
-func wait(t *testing.T, duration time.Duration, message string) {
-	t.Logf("%s - waiting for %v", message, duration)
-	time.Sleep(duration)
+	return latestFileName
 }
diff --git a/internal/storage/memory.go b/internal/storage/memory.go
index a1a7805ac..e3e0d9c0c 100644
--- a/internal/storage/memory.go
+++ b/internal/storage/memory.go
@@ -73,8 +73,7 @@ func (s *Memory) GetObject(_ context.Context, folder, filename string) ([]byte,
 	return v, nil
 }
 
-// GetObject returns the contents for the given object. If the object does not
-// exist, it returns ErrNotFound.
+// ListObjects returns the list of files in memory storage.
 func (s *Memory) ListObjects(_ context.Context, folder string) map[string][]byte {
 	s.lock.Lock()
 	defer s.lock.Unlock()
diff --git a/testing/enclient/enclient.go b/testing/enclient/enclient.go
index a8b3f463e..b55d4c775 100644
--- a/testing/enclient/enclient.go
+++ b/testing/enclient/enclient.go
@@ -30,7 +30,7 @@ import (
 
 const (
 	// httpTimeout is the maximum amount of time to wait for a response.
-	httpTimeout = 5 * time.Minute
+	httpTimeout = 30 * time.Second
 )
 
 type Interval int32

From b127f56a40286740d9c94272f3a3afb80d408369 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 10 Jun 2020 18:16:47 -0700
Subject: [PATCH 06/12] Fixing Ling errors.

---
 internal/integration/en_api.go       | 11 +++++------
 internal/integration/publish_test.go |  6 +++---
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/internal/integration/en_api.go b/internal/integration/en_api.go
index fb2700599..b6951d0bf 100644
--- a/internal/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Contains Exposure Notifications API client to support integration testing.
 package integration
 
 import (
@@ -72,22 +73,20 @@ func (server EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish
 
 func (server EnServerClient) ExportBatches(t *testing.T) {
 	var bts []byte
-	requestUrl := "/export/create-batches"
-	resp, err := server.postRequest(requestUrl, bts)
+	resp, err := server.postRequest("/export/create-batches", bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
 	log.Printf("response: %v", resp.Status)
-	t.Logf("Create batches request is sent to %v", requestUrl)
+	t.Logf("Create batches request is sent to %v", "/export/create-batches")
 }
 
 func (server EnServerClient) StartExportWorkers(t *testing.T) {
 	var bts []byte
-	requestUrl := "/export/do-work"
-	resp, err := server.postRequest(requestUrl, bts)
+	resp, err := server.postRequest("/export/do-work", bts)
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
 	log.Printf("response: %v", resp.Status)
-	t.Logf("Export worker request is sent to %v", requestUrl)
+	t.Logf("Export worker request is sent to %v", "/export/do-work")
 }
diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index d7d47b699..35af863b5 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -21,7 +21,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/protobuf/proto"
 	authorizedappmodel "github.com/google/exposure-notifications-server/internal/authorizedapp/model"
 	exportapi "github.com/google/exposure-notifications-server/internal/export"
 	exportdatabase "github.com/google/exposure-notifications-server/internal/export/database"
@@ -35,6 +34,7 @@ import (
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
+	"google.golang.org/protobuf/proto"
 )
 
 func TestPublish(t *testing.T) {
@@ -195,7 +195,7 @@ func TestPublish(t *testing.T) {
 }
 
 func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context, env *serverenv.ServerEnv, memory *storage.Memory) *export.TemporaryExposureKeyExport {
-	exportFile := getLatestFile(t, memory, ctx, exportDir)
+	exportFile := getLatestFile(memory, ctx, exportDir)
 	if exportFile == "" {
 		t.Fatalf("Can't find export files in blobstore: %v", exportDir)
 	}
@@ -215,7 +215,7 @@ func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context,
 	return keyExport
 }
 
-func getLatestFile(t *testing.T, blobstore *storage.Memory, ctx context.Context, exportDir string) string {
+func getLatestFile(blobstore *storage.Memory, ctx context.Context, exportDir string) string {
 	files := blobstore.ListObjects(ctx, exportDir)
 
 	latestFileName := ""

From c2a7907d0c81c3030f63d8052152c41002489d01 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 10 Jun 2020 18:22:25 -0700
Subject: [PATCH 07/12] Fixing Ling errors.

---
 internal/integration/en_api.go           | 2 +-
 internal/integration/integration_test.go | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/internal/integration/en_api.go b/internal/integration/en_api.go
index b6951d0bf..a554a9fa2 100644
--- a/internal/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Contains Exposure Notifications API client to support integration testing.
 package integration
 
 import (
@@ -27,6 +26,7 @@ import (
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 )
 
+// Contains Exposure Notifications API client to support integration testing.
 type EnServerClient struct {
 	client *http.Client
 }
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index e52833d63..d1000c3b3 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Package integration contains EN Server integration tests.
 package integration
 
 import (

From 40e1ef98096e03e404df2aa1736f34bb1f0d8748 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 10 Jun 2020 20:07:24 -0700
Subject: [PATCH 08/12] Making presubmit happy:

Fixing lint errors.
Fixing cmp + proto incompatibility.
---
 internal/integration/en_api.go           |  3 +-
 internal/integration/integration_test.go |  1 -
 internal/integration/publish_test.go     | 37 ++++++++++--------------
 3 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/internal/integration/en_api.go b/internal/integration/en_api.go
index a554a9fa2..a218ff08b 100644
--- a/internal/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// Package integration contains EN Server integration tests.
 package integration
 
 import (
@@ -26,7 +27,7 @@ import (
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 )
 
-// Contains Exposure Notifications API client to support integration testing.
+// EnServerClient provides Exposure Notifications API client to support integration testing.
 type EnServerClient struct {
 	client *http.Client
 }
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index d1000c3b3..e52833d63 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package integration contains EN Server integration tests.
 package integration
 
 import (
diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 35af863b5..624e98267 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -117,6 +117,7 @@ func TestPublish(t *testing.T) {
 
 	var exposures []*publishmodel.Exposure
 	if _, err := publishdb.New(db).IterateExposures(ctx, criteria, func(m *publishmodel.Exposure) error {
+		t.Logf("NEW EXPOSURE: %v", m)
 		exposures = append(exposures, m)
 		return nil
 	}); err != nil {
@@ -143,43 +144,37 @@ func TestPublish(t *testing.T) {
 
 	got := keyExport
 
-	wantedKeysMap := make(map[string]export.TemporaryExposureKey)
+	wantedKeysMap := make(map[string]*export.TemporaryExposureKey)
 	for _, key := range payload.Keys {
-		wantedKeysMap[key.Key] = export.TemporaryExposureKey{
+		wantedKeysMap[key.Key] = &export.TemporaryExposureKey{
 			KeyData:                    util.DecodeKey(key.Key),
 			TransmissionRiskLevel:      proto.Int32(int32(key.TransmissionRisk)),
 			RollingStartIntervalNumber: proto.Int32(key.IntervalNumber),
 		}
 	}
 
-	want := export.TemporaryExposureKeyExport{
-		StartTimestamp: nil,
-		EndTimestamp:   nil,
-		Region:         proto.String("TEST"),
-		BatchNum:       proto.Int32(1),
-		BatchSize:      proto.Int32(1),
-		SignatureInfos: nil,
-		Keys:           nil,
+	want := &export.TemporaryExposureKeyExport{
+		Region:    proto.String("TEST"),
+		BatchNum:  proto.Int32(1),
+		BatchSize: proto.Int32(1),
 	}
 
-	options := []cmp.Option{
-		cmpopts.IgnoreFields(want, "StartTimestamp"),
-		cmpopts.IgnoreFields(want, "EndTimestamp"),
-		cmpopts.IgnoreFields(want, "SignatureInfos"),
-		cmpopts.IgnoreFields(want, "Keys"),
-		cmpopts.IgnoreUnexported(want),
+	if !cmp.Equal(want.BatchSize, got.BatchSize) {
+		t.Errorf("Invalid BatchSize: want: %v, got: %v", want.BatchSize, got.BatchSize)
 	}
 
-	diff := cmp.Diff(got, &want, options...)
-	if diff != "" {
-		t.Errorf("%v", diff)
+	if !cmp.Equal(want.BatchNum, got.BatchNum) {
+		t.Errorf("Invalid BatchNum: want: %v, got: %v", want.BatchNum, got.BatchNum)
+	}
+
+	if !cmp.Equal(want.Region, got.Region) {
+		t.Errorf("Invalid Region: want: %v, got: %v", want.BatchSize, got.BatchSize)
 	}
 
 	for _, key := range got.Keys {
 		s := util.ToBase64(key.KeyData)
 		wantedKey := wantedKeysMap[s]
-		gotKey := *key
-		diff := cmp.Diff(wantedKey, gotKey, cmpopts.IgnoreUnexported(gotKey))
+		diff := cmp.Diff(wantedKey, key, cmpopts.IgnoreUnexported(export.TemporaryExposureKey{}))
 		if diff != "" {
 			t.Errorf("invalid key value: %v:%v", s, diff)
 		}

From 4491df8b3216bca2847c8a9f29b30adc16e825f8 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Wed, 10 Jun 2020 20:32:22 -0700
Subject: [PATCH 09/12] Fixing waiting params.

---
 internal/integration/publish_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 624e98267..7df64efc1 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -128,8 +128,8 @@ func TestPublish(t *testing.T) {
 		t.Errorf("expected %v to be %v: %#v", got, want, exposures)
 	}
 
-	t.Logf("Waiting %v before export batches", exportPeriod+500*time.Millisecond)
-	time.Sleep(exportPeriod + 500*time.Millisecond)
+	t.Logf("Waiting %v before export batches", exportPeriod+1*time.Second)
+	time.Sleep(exportPeriod + 1*time.Second)
 	enClient.ExportBatches(t)
 
 	t.Logf("Waiting %v before starting workers", 500*time.Millisecond)

From a386d6b9e48714dbdd1279934342586792d7ccbd Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Thu, 11 Jun 2020 16:58:37 -0700
Subject: [PATCH 10/12] Applying PR review comments.

---
 internal/integration/integration_test.go | 14 +++++++++++++-
 internal/integration/publish_test.go     | 15 +--------------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index e52833d63..ec21541cb 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -33,7 +33,7 @@ import (
 	"github.com/google/exposure-notifications-server/internal/storage"
 )
 
-func testServer(tb testing.TB, exportConfig *export.Config) (*serverenv.ServerEnv, *http.Client) {
+func testServer(tb testing.TB) (*serverenv.ServerEnv, *http.Client) {
 	tb.Helper()
 
 	ctx := context.Background()
@@ -96,6 +96,18 @@ func testServer(tb testing.TB, exportConfig *export.Config) (*serverenv.ServerEn
 	}
 	mux.Handle("/cleanup-exposure", cleanupExposureHandler)
 
+	// Export
+	exportConfig := &export.Config{
+		CreateTimeout:  10 * time.Second,
+		WorkerTimeout:  10 * time.Second,
+		MinRecords:     1,
+		PaddingRange:   1,
+		MaxRecords:     10000,
+		TruncateWindow: 1 * time.Millisecond,
+		MinWindowAge:   1 * time.Second,
+		TTL:            336 * time.Hour,
+	}
+
 	exportServer, err := export.NewServer(exportConfig, env)
 	if err != nil {
 		tb.Fatal(err)
diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 7df64efc1..392bf9470 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -42,19 +42,7 @@ func TestPublish(t *testing.T) {
 
 	ctx := context.Background()
 
-	// Export
-	exportConfig := &exportapi.Config{
-		CreateTimeout:  10 * time.Second,
-		WorkerTimeout:  10 * time.Second,
-		MinRecords:     1,
-		PaddingRange:   1,
-		MaxRecords:     10000,
-		TruncateWindow: 1 * time.Millisecond,
-		MinWindowAge:   1 * time.Second,
-		TTL:            336 * time.Hour,
-	}
-
-	env, client := testServer(t, exportConfig)
+	env, client := testServer(t)
 	db := env.Database()
 	enClient := &EnServerClient{client: client}
 
@@ -117,7 +105,6 @@ func TestPublish(t *testing.T) {
 
 	var exposures []*publishmodel.Exposure
 	if _, err := publishdb.New(db).IterateExposures(ctx, criteria, func(m *publishmodel.Exposure) error {
-		t.Logf("NEW EXPOSURE: %v", m)
 		exposures = append(exposures, m)
 		return nil
 	}); err != nil {

From c0d82528de5092814b482d8ebde3789abf4d01e7 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Thu, 11 Jun 2020 18:13:18 -0700
Subject: [PATCH 11/12] Applying PR review comments.

---
 internal/integration/en_api.go | 41 ++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/internal/integration/en_api.go b/internal/integration/en_api.go
index a218ff08b..242235397 100644
--- a/internal/integration/en_api.go
+++ b/internal/integration/en_api.go
@@ -32,6 +32,15 @@ type EnServerClient struct {
 	client *http.Client
 }
 
+func (server EnServerClient) getRequest(url string) (*http.Response, error) {
+	resp, err := server.client.Get(url)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	return unwrapResponse(resp)
+}
+
 // Posts requests to the specified url.
 // This methods attempts to serialize data argument as a json.
 func (server EnServerClient) postRequest(url string, data interface{}) (*http.Response, error) {
@@ -50,17 +59,7 @@ func (server EnServerClient) postRequest(url string, data interface{}) (*http.Re
 		return nil, err
 	}
 	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		// Return error upstream.
-		body, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			return nil, fmt.Errorf("failed to copy error body (%d): %w", resp.StatusCode, err)
-		}
-		return resp, fmt.Errorf("post request failed with status: %v\n%v", resp.StatusCode, body)
-	}
-
-	return resp, nil
+	return unwrapResponse(resp)
 }
 
 func (server EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish) {
@@ -73,21 +72,29 @@ func (server EnServerClient) PublishKeys(t *testing.T, request verifyapi.Publish
 }
 
 func (server EnServerClient) ExportBatches(t *testing.T) {
-	var bts []byte
-	resp, err := server.postRequest("/export/create-batches", bts)
+	resp, err := server.getRequest("/export/create-batches")
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
-	log.Printf("response: %v", resp.Status)
 	t.Logf("Create batches request is sent to %v", "/export/create-batches")
 }
 
 func (server EnServerClient) StartExportWorkers(t *testing.T) {
-	var bts []byte
-	resp, err := server.postRequest("/export/do-work", bts)
+	resp, err := server.getRequest("/export/do-work")
 	if err != nil {
 		t.Fatalf("request failed: %v, %v", err, resp)
 	}
-	log.Printf("response: %v", resp.Status)
 	t.Logf("Export worker request is sent to %v", "/export/do-work")
 }
+
+func unwrapResponse(resp *http.Response) (*http.Response, error) {
+	if resp.StatusCode != http.StatusOK {
+		// Return error upstream.
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to copy error body (%d): %w", resp.StatusCode, err)
+		}
+		return resp, fmt.Errorf("request failed with status: %v\n%v", resp.StatusCode, body)
+	}
+	return resp, nil
+}

From 6a26d70cdf0c35e1e07eaaaafeb30e9f37c43ec8 Mon Sep 17 00:00:00 2001
From: Max Gulimonov <maximg@google.com>
Date: Tue, 16 Jun 2020 16:06:40 -0700
Subject: [PATCH 12/12] Applying PR review comments.

---
 internal/integration/publish_test.go | 36 ++++++++++++++--------------
 internal/storage/memory.go           | 15 ------------
 2 files changed, 18 insertions(+), 33 deletions(-)

diff --git a/internal/integration/publish_test.go b/internal/integration/publish_test.go
index 392bf9470..bdc9bd150 100644
--- a/internal/integration/publish_test.go
+++ b/internal/integration/publish_test.go
@@ -29,7 +29,6 @@ import (
 	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/storage"
 	"github.com/google/exposure-notifications-server/internal/util"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 	"github.com/google/go-cmp/cmp"
@@ -123,11 +122,7 @@ func TestPublish(t *testing.T) {
 	time.Sleep(500 * time.Millisecond)
 	enClient.StartExportWorkers(t)
 
-	memory, ok := env.Blobstore().(*storage.Memory)
-	if !ok {
-		t.Fatalf("can't use %t blobstore for verification", env.Blobstore())
-	}
-	keyExport := getKeysFromLatestBatch(t, "my-bucket", ctx, env, memory)
+	keyExport := getKeysFromLatestBatch(t, "my-bucket", ctx, env)
 
 	got := keyExport
 
@@ -146,16 +141,16 @@ func TestPublish(t *testing.T) {
 		BatchSize: proto.Int32(1),
 	}
 
-	if !cmp.Equal(want.BatchSize, got.BatchSize) {
-		t.Errorf("Invalid BatchSize: want: %v, got: %v", want.BatchSize, got.BatchSize)
+	if *want.BatchSize != *got.BatchSize {
+		t.Errorf("Invalid BatchSize: want: %v, got: %v", *want.BatchSize, *got.BatchSize)
 	}
 
-	if !cmp.Equal(want.BatchNum, got.BatchNum) {
-		t.Errorf("Invalid BatchNum: want: %v, got: %v", want.BatchNum, got.BatchNum)
+	if *want.BatchNum != *got.BatchNum {
+		t.Errorf("Invalid BatchNum: want: %v, got: %v", *want.BatchNum, *got.BatchNum)
 	}
 
-	if !cmp.Equal(want.Region, got.Region) {
-		t.Errorf("Invalid Region: want: %v, got: %v", want.BatchSize, got.BatchSize)
+	if *want.Region != *got.Region {
+		t.Errorf("Invalid Region: want: %v, got: %v", *want.BatchSize, *got.BatchSize)
 	}
 
 	for _, key := range got.Keys {
@@ -176,15 +171,20 @@ func TestPublish(t *testing.T) {
 	// TODO: verify signature
 }
 
-func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context, env *serverenv.ServerEnv, memory *storage.Memory) *export.TemporaryExposureKeyExport {
-	exportFile := getLatestFile(memory, ctx, exportDir)
+func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context, env *serverenv.ServerEnv) *export.TemporaryExposureKeyExport {
+	readmeBlob, err := env.Blobstore().GetObject(ctx, exportDir, "index.txt")
+	if err != nil {
+		t.Fatalf("Can't file index.txt in blobstore: %v", err)
+	}
+
+	exportFile := getLatestFile(readmeBlob)
 	if exportFile == "" {
 		t.Fatalf("Can't find export files in blobstore: %v", exportDir)
 	}
 
 	t.Logf("Reading keys data from: %v", exportFile)
 
-	blob, err := env.Blobstore().GetObject(ctx, "", exportFile)
+	blob, err := env.Blobstore().GetObject(ctx, exportDir, exportFile)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -197,11 +197,11 @@ func getKeysFromLatestBatch(t *testing.T, exportDir string, ctx context.Context,
 	return keyExport
 }
 
-func getLatestFile(blobstore *storage.Memory, ctx context.Context, exportDir string) string {
-	files := blobstore.ListObjects(ctx, exportDir)
+func getLatestFile(indexBlob []byte) string {
+	files := strings.Split(string(indexBlob), "\n")
 
 	latestFileName := ""
-	for fileName := range files {
+	for _, fileName := range files {
 		if strings.HasSuffix(fileName, "zip") {
 			if latestFileName == "" {
 				latestFileName = fileName
diff --git a/internal/storage/memory.go b/internal/storage/memory.go
index e3e0d9c0c..83ea547a3 100644
--- a/internal/storage/memory.go
+++ b/internal/storage/memory.go
@@ -17,7 +17,6 @@ package storage
 import (
 	"context"
 	"path"
-	"strings"
 	"sync"
 )
 
@@ -72,17 +71,3 @@ func (s *Memory) GetObject(_ context.Context, folder, filename string) ([]byte,
 	}
 	return v, nil
 }
-
-// ListObjects returns the list of files in memory storage.
-func (s *Memory) ListObjects(_ context.Context, folder string) map[string][]byte {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	result := make(map[string][]byte)
-	for k, v := range s.data {
-		if strings.HasPrefix(k, folder+"/") {
-			result[k] = v
-		}
-	}
-	return result
-}