Repository: google/exposure-notifications-server
Pull request #722 (2020-07-15 16:01:14)

V1.5export
Add new TEK metadata fields to generated exports
Add revised keys to output
Work on #663

this is a re-open as gh wasn't picking up my merge fixes.

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/722#" title="Author self-approved">mikehelmick</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]

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":[]} -->
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/722#" title="Author self-approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/722#pullrequestreview-449184434" title="LGTM">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/internal/export/exportfile.go b/internal/export/exportfile.go
index 3150c59..22aa6ec 100644
--- a/internal/export/exportfile.go
+++ b/internal/export/exportfile.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/google/exposure-notifications-server/internal/export/model"
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
+	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 
 	"github.com/google/exposure-notifications-server/internal/pb/export"
 
@@ -51,9 +52,9 @@ type Signer struct {
 }
 
 // MarshalExportFile converts the inputs into an encoded byte array.
-func MarshalExportFile(eb *model.ExportBatch, exposures []*publishmodel.Exposure, batchNum, batchSize int, signers []*Signer) ([]byte, error) {
+func MarshalExportFile(eb *model.ExportBatch, exposures, revisedExposures []*publishmodel.Exposure, batchNum, batchSize int, signers []*Signer) ([]byte, error) {
 	// create main exposure key export binary
-	expContents, err := marshalContents(eb, exposures, int32(batchNum), int32(batchSize), signers)
+	expContents, err := marshalContents(eb, exposures, revisedExposures, int32(batchNum), int32(batchSize), signers)
 	if err != nil {
 		return nil, fmt.Errorf("unable to marshal exposure keys: %w", err)
 	}
@@ -131,7 +132,41 @@ func unmarshalContent(file *zip.File) (*export.TemporaryExposureKeyExport, error
 	return message, nil
 }
 
-func marshalContents(eb *model.ExportBatch, exposures []*publishmodel.Exposure, batchNum int32, batchSize int32, signers []*Signer) ([]byte, error) {
+func sortExposures(exposures []*publishmodel.Exposure) {
+	sort.Slice(exposures, func(i, j int) bool {
+		return bytes.Compare(exposures[i].ExposureKey, exposures[j].ExposureKey) < 0
+	})
+}
+
+func makeTEK(exp *publishmodel.Exposure) *export.TemporaryExposureKey {
+	pbek := export.TemporaryExposureKey{
+		KeyData:               exp.ExposureKey,
+		TransmissionRiskLevel: proto.Int32(int32(exp.TransmissionRisk)),
+	}
+	if exp.IntervalNumber != 0 {
+		pbek.RollingStartIntervalNumber = proto.Int32(exp.IntervalNumber)
+	}
+	if exp.IntervalCount != defaultIntervalCount {
+		pbek.RollingPeriod = proto.Int32(exp.IntervalCount)
+	}
+	return &pbek
+}
+
+func assignReportType(reportType *string, pbek *export.TemporaryExposureKey) {
+	if reportType == nil {
+		return
+	}
+	switch *reportType {
+	case verifyapi.ReportTypeConfirmed:
+		pbek.ReportType = export.TemporaryExposureKey_CONFIRMED_TEST.Enum()
+	case verifyapi.ReportTypeClinical:
+		pbek.ReportType = export.TemporaryExposureKey_CONFIRMED_CLINICAL_DIAGNOSIS.Enum()
+	case verifyapi.ReportTypeNegative:
+		pbek.ReportType = export.TemporaryExposureKey_REVOKED.Enum()
+	}
+}
+
+func marshalContents(eb *model.ExportBatch, exposures, revisedExposures []*publishmodel.Exposure, batchNum int32, batchSize int32, signers []*Signer) ([]byte, error) {
 	exportBytes := fixedHeader
 	if len(exportBytes) != fixedHeaderWidth {
 		return nil, fmt.Errorf("incorrect header length: %d", len(exportBytes))
@@ -139,22 +174,24 @@ func marshalContents(eb *model.ExportBatch, exposures []*publishmodel.Exposure,
 	// We want to scramble keys to ensure no associations, so arbitrarily sort them.
 	// This could be done at the db layer but doing it here makes it explicit that its
 	// important to the serialization
-	sort.Slice(exposures, func(i, j int) bool {
-		return bytes.Compare(exposures[i].ExposureKey, exposures[j].ExposureKey) < 0
-	})
+	sortExposures(exposures)
 	var pbeks []*export.TemporaryExposureKey
 	for _, exp := range exposures {
-		pbek := export.TemporaryExposureKey{
-			KeyData:               exp.ExposureKey,
-			TransmissionRiskLevel: proto.Int32(int32(exp.TransmissionRisk)),
-		}
-		if exp.IntervalNumber != 0 {
-			pbek.RollingStartIntervalNumber = proto.Int32(exp.IntervalNumber)
-		}
-		if exp.IntervalCount != defaultIntervalCount {
-			pbek.RollingPeriod = proto.Int32(exp.IntervalCount)
+		pbek := makeTEK(exp)
+		assignReportType(&exp.ReportType, pbek)
+		if exp.HasDaysSinceSymptomOnset() {
+			pbek.DaysSinceOnsetOfSymptoms = proto.Int32(*exp.DaysSinceSymptomOnset)
 		}
-		pbeks = append(pbeks, &pbek)
+		pbeks = append(pbeks, pbek)
+	}
+
+	sortExposures(revisedExposures)
+	var pbRevisedKeys []*export.TemporaryExposureKey
+	for _, exp := range revisedExposures {
+		pbek := makeTEK(exp)
+		assignReportType(exp.RevisedReportType, pbek)
+		pbek.DaysSinceOnsetOfSymptoms = exp.RevisedDaysSinceSymptomOnset
+		pbRevisedKeys = append(pbRevisedKeys, pbek)
 	}
 
 	var exportSigInfos []*export.SignatureInfo
@@ -169,6 +206,7 @@ func marshalContents(eb *model.ExportBatch, exposures []*publishmodel.Exposure,
 		BatchNum:       proto.Int32(batchNum),
 		BatchSize:      proto.Int32(batchSize),
 		Keys:           pbeks,
+		RevisedKeys:    pbRevisedKeys,
 		SignatureInfos: exportSigInfos,
 	}
 	protoBytes, err := proto.Marshal(&pbeke)
diff --git a/internal/export/worker.go b/internal/export/worker.go
index 45ce239..747ed2d 100644
--- a/internal/export/worker.go
+++ b/internal/export/worker.go
@@ -33,7 +33,6 @@ import (
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
 
 	"github.com/google/exposure-notifications-server/internal/logging"
-	"github.com/google/exposure-notifications-server/internal/util"
 
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 )
@@ -97,50 +96,85 @@ func (s *Server) handleDoWork(ctx context.Context) http.HandlerFunc {
 	}
 }
 
+type group struct {
+	exposures []*publishmodel.Exposure
+	revised   []*publishmodel.Exposure
+}
+
+func (g *group) Length() int {
+	return len(g.exposures) + len(g.revised)
+}
+
 func (s *Server) exportBatch(ctx context.Context, eb *model.ExportBatch, emitIndexForEmptyBatch bool) error {
 	logger := logging.FromContext(ctx)
 	db := s.env.Database()
 
 	logger.Infof("Processing export batch %d (root: %q, region: %s), max records per file %d", eb.BatchID, eb.FilenameRoot, eb.OutputRegion, s.config.MaxRecords)
 
+	// Criteria starts w/ non-revised keys.
+	// Will be changed later to grab the revised keys.
 	criteria := publishdatabase.IterateExposuresCriteria{
 		SinceTimestamp:      eb.StartTimestamp,
 		UntilTimestamp:      eb.EndTimestamp,
 		IncludeRegions:      eb.EffectiveInputRegions(),
 		OnlyLocalProvenance: false, // include federated ids
+		RevisedKeys:         false,
 	}
 
 	// Build up groups of exposures in memory. We need to use memory so we can
 	// determine the total number of groups (which is embedded in each export
 	// file). This technique avoids SELECT COUNT which would lock the database
 	// slowing new uploads.
-	var groups [][]*publishmodel.Exposure
-	var exposures []*publishmodel.Exposure
+	groups := []*group{}
+	nextGroup := group{}
+	totalNewKeys, totalRevisedKeys := 0, 0
 
 	_, err := publishdatabase.New(db).IterateExposures(ctx, criteria, func(exp *publishmodel.Exposure) error {
-		exposures = append(exposures, exp)
-		if len(exposures) == s.config.MaxRecords {
-			groups = append(groups, exposures)
-			exposures = nil
+		nextGroup.exposures = append(nextGroup.exposures, exp)
+		totalNewKeys++
+		if nextGroup.Length() == s.config.MaxRecords {
+			groups = append(groups, &nextGroup)
+			nextGroup = group{}
 		}
 		return nil
 	})
-
 	if err != nil {
 		return fmt.Errorf("iterating exposures: %w", err)
 	}
-	// Create a group for any remaining keys.
-	if len(exposures) > 0 {
-		groups = append(groups, exposures)
+
+	// go get the revised keys.
+	criteria.RevisedKeys = true
+	_, err = publishdatabase.New(db).IterateExposures(ctx, criteria, func(exp *publishmodel.Exposure) error {
+		nextGroup.revised = append(nextGroup.revised, exp)
+		totalRevisedKeys++
+		if nextGroup.Length() == s.config.MaxRecords {
+			groups = append(groups, &nextGroup)
+			nextGroup = group{}
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("iterating revised exposures: %w", err)
+	}
+
+	// If the last group has anything, add it to the list.
+	if nextGroup.Length() > 0 {
+		groups = append(groups, &nextGroup)
 	}
 
 	if len(groups) == 0 {
 		logger.Infof("No records for export batch %d", eb.BatchID)
-	}
+	} else {
 
-	exposures, err = ensureMinNumExposures(exposures, eb.OutputRegion, s.config.MinRecords, s.config.PaddingRange)
-	if err != nil {
-		return fmt.Errorf("ensureMinNumExposures: %w", err)
+		lastGroup := groups[len(groups)-1]
+		lastGroup.exposures, err = ensureMinNumExposures(lastGroup.exposures, eb.OutputRegion, s.config.MinRecords, s.config.PaddingRange)
+		if err != nil {
+			return fmt.Errorf("ensureMinNumExposures: %w", err)
+		}
+		lastGroup.revised, err = ensureMinNumExposures(lastGroup.revised, eb.OutputRegion, s.config.MinRecords, s.config.PaddingRange)
+		if err != nil {
+			return fmt.Errorf("ensureMinNumExposures: %w", err)
+		}
 	}
 
 	// Load the non-expired signature infos associated with this export batch.
@@ -152,7 +186,7 @@ func (s *Server) exportBatch(ctx context.Context, eb *model.ExportBatch, emitInd
 	// Create the export files.
 	batchSize := len(groups)
 	var objectNames []string
-	for i, exposures := range groups {
+	for i, group := range groups {
 		if ctx.Err() != nil {
 			logger.Infof("Timed out writing export files for batch %s, the entire batch will be retried once the batch lease expires on %v", eb.BatchID, eb.LeaseExpires)
 			return nil
@@ -162,11 +196,12 @@ func (s *Server) exportBatch(ctx context.Context, eb *model.ExportBatch, emitInd
 		// use of network.
 		objectName, err := s.createFile(ctx,
 			createFileInfo{
-				exposures:      exposures,
-				exportBatch:    eb,
-				signatureInfos: sigInfos,
-				batchNum:       i + 1,
-				batchSize:      batchSize,
+				exposures:        group.exposures,
+				revisedExposures: group.revised,
+				exportBatch:      eb,
+				signatureInfos:   sigInfos,
+				batchNum:         i + 1,
+				batchSize:        batchSize,
 			})
 		if err != nil {
 			return fmt.Errorf("creating export file %d for batch %d: %w", i+1, eb.BatchID, err)
@@ -191,11 +226,12 @@ func (s *Server) exportBatch(ctx context.Context, eb *model.ExportBatch, emitInd
 }
 
 type createFileInfo struct {
-	exposures      []*publishmodel.Exposure
-	exportBatch    *model.ExportBatch
-	signatureInfos []*model.SignatureInfo
-	batchNum       int
-	batchSize      int
+	exposures        []*publishmodel.Exposure
+	revisedExposures []*publishmodel.Exposure
+	exportBatch      *model.ExportBatch
+	signatureInfos   []*model.SignatureInfo
+	batchNum         int
+	batchSize        int
 }
 
 func (s *Server) createFile(ctx context.Context, cfi createFileInfo) (string, error) {
@@ -211,7 +247,7 @@ func (s *Server) createFile(ctx context.Context, cfi createFileInfo) (string, er
 	}
 
 	// Generate exposure key export file.
-	data, err := MarshalExportFile(cfi.exportBatch, cfi.exposures, cfi.batchNum, cfi.batchSize, signers)
+	data, err := MarshalExportFile(cfi.exportBatch, cfi.exposures, cfi.revisedExposures, cfi.batchNum, cfi.batchSize, signers)
 	if err != nil {
 		return "", fmt.Errorf("marshaling export file: %w", err)
 	}
@@ -326,6 +362,7 @@ func ensureMinNumExposures(exposures []*publishmodel.Exposure, region string, mi
 
 	extra, _ := randomInt(0, jitter)
 	target := minLength + extra
+	sourceLen := len(exposures) - 1
 
 	for len(exposures) < target {
 		// Pieces needed are
@@ -337,34 +374,23 @@ func ensureMinNumExposures(exposures []*publishmodel.Exposure, region string, mi
 			return nil, fmt.Errorf("rand.Read: %w", err)
 		}
 
-		// Transmission risk is within the bounds.
-		transmissionRisk, err := util.RandomTransmissionRisk()
-		if err != nil {
-			return nil, fmt.Errorf("randomInt: %w", err)
-		}
-
-		// The interval number is pulled from an existing one in the batch
-		// at random.
-		fromIdx, err := randomInt(0, len(exposures)-1)
-		if err != nil {
-			return nil, fmt.Errorf("randomInt: %w", err)
-		}
-		intervalNumber := exposures[fromIdx].IntervalNumber
-		// Same with the interval count.
-		fromIdx, err = randomInt(0, len(exposures)-1)
+		// Copy timing and report data from a key.
+		fromIdx, err := randomInt(0, sourceLen)
 		if err != nil {
 			return nil, fmt.Errorf("randomInt: %w", err)
 		}
-		intervalCount := exposures[fromIdx].IntervalCount
 
 		ek := &publishmodel.Exposure{
-			ExposureKey:      eKey,
-			TransmissionRisk: transmissionRisk,
-			Regions:          []string{region},
-			IntervalNumber:   intervalNumber,
-			IntervalCount:    intervalCount,
-			// The rest of the publishmodel.Exposure fields are not used in the export
-			// file.
+			ExposureKey:                  eKey,
+			TransmissionRisk:             exposures[fromIdx].TransmissionRisk,
+			Regions:                      []string{region},
+			IntervalNumber:               exposures[fromIdx].IntervalNumber,
+			IntervalCount:                exposures[fromIdx].IntervalCount,
+			ReportType:                   exposures[fromIdx].ReportType,
+			DaysSinceSymptomOnset:        exposures[fromIdx].DaysSinceSymptomOnset,
+			RevisedReportType:            exposures[fromIdx].RevisedReportType,
+			RevisedDaysSinceSymptomOnset: exposures[fromIdx].RevisedDaysSinceSymptomOnset,
+			// The rest of the publishmodel.Exposure fields are not used in the export file.
 		}
 		exposures = append(exposures, ek)
 	}
diff --git a/internal/publish/model/exposure_model.go b/internal/publish/model/exposure_model.go
index b6d4669..7ce5c7d 100644
--- a/internal/publish/model/exposure_model.go
+++ b/internal/publish/model/exposure_model.go
@@ -32,10 +32,12 @@ import (
 )
 
 var (
-	ErrorExposureKeyMismatch     = fmt.Errorf("attempted to revise a key with a different key")
-	ErrorHealthAuthorityMismatch = fmt.Errorf("key revisition attempted by different health authority")
-	ErrorNonLocalProvenance      = fmt.Errorf("key not origionally uploaded to this server, cannot revise")
-	ErrorKeyAlreadyRevised       = fmt.Errorf("key has already been revised and cannot be revised again")
+	// ErrorExposureKeyMismatch - internal coding error, tried to revise key A by passing in key B
+	ErrorExposureKeyMismatch = fmt.Errorf("attempted to revise a key with a different key")
+	// ErrorNonLocalProvenance - key revesion attempted on federated key, which is not allowed
+	ErrorNonLocalProvenance = fmt.Errorf("key not origionally uploaded to this server, cannot revise")
+	// ErrorKeyAlreadyRevised - attempt to revise a key that has already been revised.
+	ErrorKeyAlreadyRevised = fmt.Errorf("key has already been revised and cannot be revised again")
 )
 
 // Exposure represents the record as stored in the database
diff --git a/tools/export-generate/main.go b/tools/export-generate/main.go
index 444c334..2217fe7 100644
--- a/tools/export-generate/main.go
+++ b/tools/export-generate/main.go
@@ -189,6 +189,7 @@ func main() {
 type exportFileWriter struct {
 	exportBatch *model.ExportBatch
 	exposures   []*publishmodel.Exposure
+	revisions   []*publishmodel.Exposure
 	curBatch    int
 	numBatches  int
 	totalKeys   int
@@ -204,7 +205,7 @@ func (e *exportFileWriter) writeFile() {
 		SignatureInfo: signatureInfo,
 		Signer:        e.privateKey,
 	}
-	data, err := export.MarshalExportFile(e.exportBatch, e.exposures, e.curBatch, e.numBatches, []*export.Signer{signer})
+	data, err := export.MarshalExportFile(e.exportBatch, e.exposures, e.revisions, e.curBatch, e.numBatches, []*export.Signer{signer})
 	if err != nil {
 		log.Fatalf("error marshaling export file: %v", err)
 	}


Test patch:
diff --git a/internal/export/exportfile_test.go b/internal/export/exportfile_test.go
index 0d8a7e1..1dfbcfa 100644
--- a/internal/export/exportfile_test.go
+++ b/internal/export/exportfile_test.go
@@ -23,6 +23,8 @@ import (
 	"github.com/google/exposure-notifications-server/internal/export/model"
 	"github.com/google/exposure-notifications-server/internal/pb/export"
 	publishmodel "github.com/google/exposure-notifications-server/internal/publish/model"
+	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"
@@ -47,22 +49,39 @@ func TestMarshalUnmarshalExportFile(t *testing.T) {
 
 	exposures := []*publishmodel.Exposure{
 		{
-			ExposureKey:      []byte("ABC"),
-			Regions:          []string{"US"},
-			IntervalNumber:   18,
-			IntervalCount:    0,
-			CreatedAt:        batchStartTime,
-			LocalProvenance:  true,
-			TransmissionRisk: 8,
+			ExposureKey:           []byte("ABC"),
+			Regions:               []string{"US"},
+			IntervalNumber:        18,
+			IntervalCount:         0,
+			CreatedAt:             batchStartTime,
+			LocalProvenance:       true,
+			TransmissionRisk:      8,
+			DaysSinceSymptomOnset: proto.Int32(0),
+			ReportType:            verifyapi.ReportTypeClinical,
+		},
+		{
+			ExposureKey:           []byte("DEF"),
+			Regions:               []string{"CA"},
+			IntervalNumber:        118,
+			IntervalCount:         1,
+			CreatedAt:             batchEndTime,
+			LocalProvenance:       true,
+			TransmissionRisk:      1,
+			DaysSinceSymptomOnset: proto.Int32(-1),
+			ReportType:            verifyapi.ReportTypeConfirmed,
 		},
+	}
+	revisedExposures := []*publishmodel.Exposure{
 		{
-			ExposureKey:      []byte("DEF"),
-			Regions:          []string{"CA"},
-			IntervalNumber:   118,
-			IntervalCount:    1,
-			CreatedAt:        batchEndTime,
-			LocalProvenance:  true,
-			TransmissionRisk: 1,
+			ExposureKey:           []byte("123"),
+			Regions:               []string{"US"},
+			IntervalNumber:        100,
+			IntervalCount:         144,
+			CreatedAt:             batchStartTime,
+			LocalProvenance:       true,
+			TransmissionRisk:      4,
+			DaysSinceSymptomOnset: proto.Int32(2),
+			RevisedReportType:     proto.String(verifyapi.ReportTypeNegative),
 		},
 	}
 
@@ -75,7 +94,7 @@ func TestMarshalUnmarshalExportFile(t *testing.T) {
 
 	signer := &customTestSigner{}
 
-	blob, err := MarshalExportFile(batch, exposures, 1, 1, []*Signer{
+	blob, err := MarshalExportFile(batch, exposures, revisedExposures, 1, 1, []*Signer{
 		{SignatureInfo: signatureInfo, Signer: signer},
 	})
 	if err != nil {
@@ -101,12 +120,24 @@ func TestMarshalUnmarshalExportFile(t *testing.T) {
 			TransmissionRiskLevel:      proto.Int32(8),
 			RollingStartIntervalNumber: proto.Int32(18),
 			RollingPeriod:              proto.Int32(0),
+			ReportType:                 export.TemporaryExposureKey_CONFIRMED_CLINICAL_DIAGNOSIS.Enum(),
+			DaysSinceOnsetOfSymptoms:   proto.Int32(0),
 		},
 		{
 			KeyData:                    []byte("DEF"),
 			TransmissionRiskLevel:      proto.Int32(1),
 			RollingStartIntervalNumber: proto.Int32(118),
 			RollingPeriod:              proto.Int32(1),
+			ReportType:                 export.TemporaryExposureKey_CONFIRMED_TEST.Enum(),
+			DaysSinceOnsetOfSymptoms:   proto.Int32(-1),
+		},
+	}
+	revisedKeys := []*export.TemporaryExposureKey{
+		{
+			KeyData:                    []byte("123"),
+			TransmissionRiskLevel:      proto.Int32(4),
+			RollingStartIntervalNumber: proto.Int32(100),
+			ReportType:                 export.TemporaryExposureKey_REVOKED.Enum(),
 		},
 	}
 
@@ -118,6 +149,7 @@ func TestMarshalUnmarshalExportFile(t *testing.T) {
 		BatchSize:      proto.Int32(1),
 		SignatureInfos: infos,
 		Keys:           keys,
+		RevisedKeys:    revisedKeys,
 	}
 
 	ignoredTemporaryExposureKeyExportFields := cmpopts.IgnoreUnexported(export.TemporaryExposureKeyExport{})
diff --git a/internal/export/worker_test.go b/internal/export/worker_test.go
index cb89139..9d0d59c 100644
--- a/internal/export/worker_test.go
+++ b/internal/export/worker_test.go
@@ -15,12 +15,11 @@
 package export
 
 import (
-	"crypto/rand"
-	"fmt"
 	"testing"
 
 	"github.com/google/exposure-notifications-server/internal/publish/model"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
+	"google.golang.org/protobuf/proto"
 )
 
 func TestRandomInt(t *testing.T) {
@@ -57,39 +56,41 @@ func TestDoNotPadZeroLength(t *testing.T) {
 	}
 }
 
-func addExposure(t *testing.T, exposures []*model.Exposure, interval, count int32, risk int) []*model.Exposure {
-	key := make([]byte, verifyapi.KeyLength)
-	_, err := rand.Read(key)
-	if err != nil {
-		t.Fatalf("error getting random data: %v", err)
-	}
-	return append(exposures,
-		&model.Exposure{
-			ExposureKey:      key,
-			TransmissionRisk: risk,
-			IntervalNumber:   interval,
-			IntervalCount:    count,
-		})
-}
-
 func TestEnsureMinExposures(t *testing.T) {
-	expectedTR := make(map[int]bool)
-	for i := verifyapi.MinTransmissionRisk; i <= verifyapi.MaxTransmissionRisk; i++ {
-		expectedTR[i] = true
-	}
 	// Insert a few exposures - that will be used to base the interval information off of.
-	exposures := make([]*model.Exposure, 0)
-	exposures = addExposure(t, exposures, 123456, 144, 0)
-	exposures = addExposure(t, exposures, 789101, 88, 0)
-	// all of these combinations are expected
-	eIntervals := make(map[string]bool)
-	eIntervals["123456.144"] = true // covered by input
-	eIntervals["123456.88"] = false
-	eIntervals["789101.144"] = false
-	eIntervals["789101.88"] = true
+	exposures := []*model.Exposure{
+		{
+			ExposureKey:           []byte{1},
+			TransmissionRisk:      verifyapi.TransmissionRiskConfirmedStandard,
+			IntervalNumber:        123456,
+			IntervalCount:         144,
+			DaysSinceSymptomOnset: proto.Int32(0),
+			ReportType:            verifyapi.ReportTypeConfirmed,
+		},
+		{
+			ExposureKey:           []byte{2},
+			TransmissionRisk:      verifyapi.TransmissionRiskConfirmedStandard,
+			IntervalNumber:        123456 + 144,
+			IntervalCount:         144,
+			DaysSinceSymptomOnset: proto.Int32(1),
+			ReportType:            verifyapi.ReportTypeConfirmed,
+		},
+		{
+			ExposureKey:           []byte{3},
+			TransmissionRisk:      verifyapi.TransmissionRiskConfirmedStandard,
+			IntervalNumber:        123456 - 144,
+			IntervalCount:         144,
+			DaysSinceSymptomOnset: proto.Int32(-1),
+			ReportType:            verifyapi.ReportTypeConfirmed,
+		},
+	}
 
-	numKeys := 4000
-	variance := 100
+	numKeys := 100
+	variance := 10
+	m := make(map[int32]int)
+	m[123456-144] = 1
+	m[123456] = 1
+	m[123456+144] = 1
 
 	// pad the download.
 	exposures, err := ensureMinNumExposures(exposures, "US", numKeys, variance)
@@ -101,18 +102,11 @@ func TestEnsureMinExposures(t *testing.T) {
 	}
 
 	for _, e := range exposures {
-		delete(expectedTR, e.TransmissionRisk)
-		eIntervals[fmt.Sprintf("%v.%v", e.IntervalNumber, e.IntervalCount)] = true
-	}
-	if len(expectedTR) != 0 {
-		t.Errorf("Didn't cover all expected transmission risks in batch of 1000: %v", expectedTR)
-	}
-	if got, want := len(eIntervals), 4; got != want {
-		t.Errorf("Unexpected number of intervalNum/count combinations, got %d, want %d", got, want)
+		m[e.IntervalNumber] = m[e.IntervalNumber] + 1
 	}
-	for k, v := range eIntervals {
-		if !v {
-			t.Errorf("interval %v was not found in output", k)
+	for k, v := range m {
+		if v < 20 {
+			t.Errorf("distribution not random, expected >= 30 keys with start interval %v, got %v", k, v)
 		}
 	}
 }