test_plan_generator: Start writing new TestUnit fields to output

We'll dual-write for now to the new and old places, then migrate recipes
to use the new fields, then drop writing the old test units.

In the mean time, a lot of this code (and tests) look bloated, but it'll
be easy to clean that up afterward.

There's an outstanding issue of figuring out if the builds actually have
the critical field set at all. This CL assumes that the field is set,
and defaults to critical := true anyway, so there's no harm.

BUG=chromium:963016
TEST=unit tests, local run

Change-Id: Ic191f1ad43bbf5218a3bdefe46ba289ae3be3123
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/infra/test_planner/+/1613340
Reviewed-by: Evan Hernandez <evanhernandez@chromium.org>
Commit-Queue: Sean Abraham <seanabraham@chromium.org>
Tested-by: Sean Abraham <seanabraham@chromium.org>
diff --git a/deploy_cipd.json b/deploy_cipd.json
index 01ae0b6..cc680de 100644
--- a/deploy_cipd.json
+++ b/deploy_cipd.json
@@ -1,6 +1,6 @@
 {
   "result": {
     "package": "chromiumos/infra/test_planner",
-    "instance_id": "LqUVQPnbis7SD1U0ZGW6cUZ1dj5A9hdF_zwc2r9Qw9wC"
+    "instance_id": "k6R1yB6jhcFt51tF1ymGiFC6lTH1huyLQ4B4pW8x4XUC"
   }
 }
\ No newline at end of file
diff --git a/src/testplans/cmd/test_plan_generator/main.go b/src/testplans/cmd/test_plan_generator/main.go
index 3ea8c69..b2145b8 100644
--- a/src/testplans/cmd/test_plan_generator/main.go
+++ b/src/testplans/cmd/test_plan_generator/main.go
@@ -34,7 +34,7 @@
 )
 
 var (
-	unmarshaler = jsonpb.Unmarshaler{AllowUnknownFields:true}
+	unmarshaler = jsonpb.Unmarshaler{AllowUnknownFields: true}
 )
 
 func cmdGenTestPlan(authOpts auth.Options) *subcommands.Command {
diff --git a/src/testplans/go.mod b/src/testplans/go.mod
index 32ece99..293b8e5 100644
--- a/src/testplans/go.mod
+++ b/src/testplans/go.mod
@@ -15,7 +15,7 @@
 	github.com/smartystreets/assertions v0.0.0-20190401211740-f487f9de1cd3 // indirect
 	github.com/smartystreets/goconvey v0.0.0-20190330032615-68dc04aab96a // indirect
 	github.com/texttheater/golang-levenshtein v0.0.0-20180516184445-d188e65d659e // indirect
-	go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190506153826-232eeaaa29a9
+	go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514225112-82c9d54cbe67
 	go.chromium.org/luci v0.0.0-20190404171609-14fb4fbce8ea
 	golang.org/x/oauth2 v0.0.0-20190402181905-9f3314589c9a // indirect
 	google.golang.org/grpc v1.19.1 // indirect
diff --git a/src/testplans/go.sum b/src/testplans/go.sum
index f068f0a..77326f8 100644
--- a/src/testplans/go.sum
+++ b/src/testplans/go.sum
@@ -117,8 +117,17 @@
 github.com/tarm/serial v0.0.0-20180830185346-98f6abe2eb07/go.mod h1:kDXzergiv9cbyO7IOYJZWg1U88JhDg3PB6klq9Hg2pA=
 github.com/texttheater/golang-levenshtein v0.0.0-20180516184445-d188e65d659e h1:T5PdfK/M1xyrHwynxMIVMWLS7f/qHwfslZphxtGnw7s=
 github.com/texttheater/golang-levenshtein v0.0.0-20180516184445-d188e65d659e/go.mod h1:XDKHRm5ThF8YJjx001LtgelzsoaEcvnA7lVWz9EeX3g=
+go.chromium.org/chromiumos/infra/proto v0.0.0-20190514195506-9aeeb525f8a2 h1:gljjPuw1pRc2lK9PLb6VvGvXgNDOuKCEmHnKlhxHUcg=
+go.chromium.org/chromiumos/infra/proto v0.0.0-20190514223940-f962e75b3f92 h1:qLY0UEyKr26Er3Ilv+7RuRgj5SGWtChdLFztrOxjB6g=
+go.chromium.org/chromiumos/infra/proto v0.0.0-20190514225112-82c9d54cbe67 h1:oF/0U7DfnhRgov+m3xk3u8KC9yrjuKP+VTZ0CWkOOkw=
 go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190506153826-232eeaaa29a9 h1:87ML72wybRl6dKe4qyvTH5Y6ihvrTxT13+xTsz7F6tE=
 go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190506153826-232eeaaa29a9/go.mod h1:Y1AiM5bpNKKecKy3wgiBwF9q1PJ/f57z86DpkNrMehk=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514195506-9aeeb525f8a2 h1:kAJWnYFu1Vvk3RdZLKz9rDDBsi4m8WPDXv7r9plSEaY=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514195506-9aeeb525f8a2/go.mod h1:Y1AiM5bpNKKecKy3wgiBwF9q1PJ/f57z86DpkNrMehk=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514223940-f962e75b3f92 h1:Ne3yU/li76ojZl+yR82ZKbVubPVigF8ybZydSe5/qjw=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514223940-f962e75b3f92/go.mod h1:Y1AiM5bpNKKecKy3wgiBwF9q1PJ/f57z86DpkNrMehk=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514225112-82c9d54cbe67 h1:gTOb9jvLQwIgx+fRhTV1TIUwWoB5SMBdvsZG+0R6/oA=
+go.chromium.org/chromiumos/infra/proto/go v0.0.0-20190514225112-82c9d54cbe67/go.mod h1:Y1AiM5bpNKKecKy3wgiBwF9q1PJ/f57z86DpkNrMehk=
 go.chromium.org/luci v0.0.0-20190404171609-14fb4fbce8ea h1:qf3Ri7sv3yJPL4E0iq4hBYn3AFSf2dNNJpuZDBXAT2E=
 go.chromium.org/luci v0.0.0-20190404171609-14fb4fbce8ea/go.mod h1:MIQewVTLvOvc0UioV0JNqTNO/RspKFS0XEeoKrOxsdM=
 go.opencensus.io v0.18.0/go.mod h1:vKdFvxhtzZ9onBp9VKHK8z/sRpBMnKAsufL7wlDrCOA=
diff --git a/src/testplans/internal/generator/generator.go b/src/testplans/internal/generator/generator.go
index f384646..ea3ec1f 100644
--- a/src/testplans/internal/generator/generator.go
+++ b/src/testplans/internal/generator/generator.go
@@ -7,13 +7,13 @@
 	"errors"
 	"fmt"
 	"github.com/golang/protobuf/proto"
+	"github.com/golang/protobuf/ptypes/wrappers"
 	"go.chromium.org/chromiumos/infra/proto/go/chromiumos"
+	"go.chromium.org/chromiumos/infra/proto/go/testplans"
+	bbproto "go.chromium.org/luci/buildbucket/proto"
 	"log"
 	"strings"
 	"testplans/internal/git"
-
-	"go.chromium.org/chromiumos/infra/proto/go/testplans"
-	bbproto "go.chromium.org/luci/buildbucket/proto"
 )
 
 type testType int
@@ -88,12 +88,7 @@
 		targetBuildResults = append(targetBuildResults, *tbr)
 	}
 
-	// The final list of TestUnits needed for the test plan.
-	testUnits, err := createTestUnits(targetBuildResults, skippableTests)
-	if err != nil {
-		return testPlan, err
-	}
-	return &testplans.GenerateTestPlanResponse{TestUnit: testUnits}, nil
+	return createResponse(targetBuildResults, skippableTests)
 }
 
 // getBuildTarget returns the build target from the given build, or empty string if none is found.
@@ -109,12 +104,12 @@
 	return bt.GetStringValue()
 }
 
-// createTestUnits creates the final list of tests required for the GenerateTestPlanResponse.
-func createTestUnits(
+// createResponse creates the final GenerateTestPlanResponse.
+func createResponse(
 	targetBuildResults []targetBuildResult,
-	skippableTests map[BuildTarget]map[testType]bool) ([]*testplans.TestUnit, error) {
+	skippableTests map[BuildTarget]map[testType]bool) (*testplans.GenerateTestPlanResponse, error) {
 
-	testUnits := make([]*testplans.TestUnit, 0)
+	resp := &testplans.GenerateTestPlanResponse{}
 targetLoop:
 	for _, tbr := range targetBuildResults {
 		pointlessBuild, ok := tbr.buildReport.Output.Properties.Fields["pointless_build"]
@@ -141,51 +136,94 @@
 		}
 		pttr := tbr.perTargetTestReqs
 		bt := chromiumos.BuildTarget{Name: string(tbr.buildTarget)}
+		tuc := &testplans.TestUnitCommon{BuildTarget: &bt, BuildPayload: bp}
+		critical := &wrappers.BoolValue{Value: tbr.buildReport.Critical != bbproto.Trinary_NO}
 		if pttr.GceTestCfg != nil {
-			testUnits = append(testUnits,
+			for _, gce := range pttr.GceTestCfg.GceTest {
+				gce.Common = withCritical(gce.Common, critical)
+			}
+			resp.TestUnit = append(resp.TestUnit,
 				&testplans.TestUnit{
 					BuildTarget:  &bt,
 					BuildPayload: bp,
 					TestCfg:      &testplans.TestUnit_GceTestCfg{GceTestCfg: pttr.GceTestCfg}})
+			resp.GceTestUnits = append(resp.GceTestUnits, &testplans.GceTestUnit{
+				Common:     tuc,
+				GceTestCfg: pttr.GceTestCfg})
 		}
 		if pttr.HwTestCfg != nil {
 			if skippableTests[tbr.buildTarget][hw] {
 				log.Printf("No HW testing needed for %s", tbr.buildTarget)
 			} else {
-				testUnits = append(testUnits,
+				for _, hw := range pttr.HwTestCfg.HwTest {
+					hw.Common = withCritical(hw.Common, critical)
+				}
+				resp.TestUnit = append(resp.TestUnit,
 					&testplans.TestUnit{
 						BuildTarget:  &bt,
 						BuildPayload: bp,
 						TestCfg:      &testplans.TestUnit_HwTestCfg{HwTestCfg: pttr.HwTestCfg}})
+				resp.HwTestUnits = append(resp.HwTestUnits, &testplans.HwTestUnit{
+					Common:    tuc,
+					HwTestCfg: pttr.HwTestCfg})
 			}
 		}
 		if pttr.MoblabVmTestCfg != nil {
-			testUnits = append(testUnits,
+			for _, moblab := range pttr.MoblabVmTestCfg.MoblabTest {
+				moblab.Common = withCritical(moblab.Common, critical)
+			}
+			resp.TestUnit = append(resp.TestUnit,
 				&testplans.TestUnit{
 					BuildTarget:  &bt,
 					BuildPayload: bp,
 					TestCfg:      &testplans.TestUnit_MoblabVmTestCfg{MoblabVmTestCfg: pttr.MoblabVmTestCfg}})
+			resp.MoblabVmTestUnits = append(resp.MoblabVmTestUnits, &testplans.MoblabVmTestUnit{
+				Common:          tuc,
+				MoblabVmTestCfg: pttr.MoblabVmTestCfg})
 		}
 		if pttr.TastVmTestCfg != nil {
-			testUnits = append(testUnits,
+			for _, tastVm := range pttr.TastVmTestCfg.TastVmTest {
+				tastVm.Common = withCritical(tastVm.Common, critical)
+			}
+			resp.TestUnit = append(resp.TestUnit,
 				&testplans.TestUnit{
 					BuildTarget:  &bt,
 					BuildPayload: bp,
 					TestCfg:      &testplans.TestUnit_TastVmTestCfg{TastVmTestCfg: pttr.TastVmTestCfg}})
+			resp.TastVmTestUnits = append(resp.TastVmTestUnits, &testplans.TastVmTestUnit{
+				Common:        tuc,
+				TastVmTestCfg: pttr.TastVmTestCfg})
 		}
 		if pttr.VmTestCfg != nil {
 			if skippableTests[tbr.buildTarget][vm] {
 				log.Printf("No VM testing needed for %s", tbr.buildTarget)
 			} else {
-				testUnits = append(testUnits,
+				for _, vm := range pttr.VmTestCfg.VmTest {
+					vm.Common = withCritical(vm.Common, critical)
+				}
+				resp.TestUnit = append(resp.TestUnit,
 					&testplans.TestUnit{
 						BuildTarget:  &bt,
 						BuildPayload: bp,
 						TestCfg:      &testplans.TestUnit_VmTestCfg{VmTestCfg: pttr.VmTestCfg}})
+				resp.VmTestUnits = append(resp.VmTestUnits, &testplans.VmTestUnit{
+					Common:    tuc,
+					VmTestCfg: pttr.VmTestCfg})
 			}
 		}
 	}
-	return testUnits, nil
+	return resp, nil
+}
+
+func withCritical(tsc *testplans.TestSuiteCommon, critical *wrappers.BoolValue) *testplans.TestSuiteCommon {
+	if tsc == nil {
+		tsc = &testplans.TestSuiteCommon{}
+	}
+	tsc.Critical = critical
+	if !critical.Value {
+		log.Printf("Marking %s as not critical", tsc.DisplayName)
+	}
+	return tsc
 }
 
 // extractSkippableTests maps BuildTargets to the test types that can be skipped for those targets,
diff --git a/src/testplans/internal/generator/generator_test.go b/src/testplans/internal/generator/generator_test.go
index 82c1cc9..dbd279e 100644
--- a/src/testplans/internal/generator/generator_test.go
+++ b/src/testplans/internal/generator/generator_test.go
@@ -4,14 +4,15 @@
 package generator
 
 import (
-	"go.chromium.org/chromiumos/infra/proto/go/chromiumos"
-	"testing"
-	"testplans/internal/git"
-
 	_struct "github.com/golang/protobuf/ptypes/struct"
+	"github.com/golang/protobuf/ptypes/wrappers"
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
+	"go.chromium.org/chromiumos/infra/proto/go/chromiumos"
 	"go.chromium.org/chromiumos/infra/proto/go/testplans"
 	bbproto "go.chromium.org/luci/buildbucket/proto"
+	"testing"
+	"testplans/internal/git"
 )
 
 const (
@@ -19,9 +20,16 @@
 	GS_PATH_PREFIX = "gs/path/"
 )
 
-func makeBuildbucketBuild(buildTarget string, status bbproto.Status, changes []*bbproto.GerritChange) *bbproto.Build {
+func makeBuildbucketBuild(buildTarget string, status bbproto.Status, changes []*bbproto.GerritChange, critical bool) *bbproto.Build {
+	var criticalVal bbproto.Trinary
+	if critical {
+		criticalVal = bbproto.Trinary_YES
+	} else {
+		criticalVal = bbproto.Trinary_NO
+	}
 	b := &bbproto.Build{
-		Input: &bbproto.Build_Input{},
+		Critical: criticalVal,
+		Input:    &bbproto.Build_Input{},
 		Output: &bbproto.Build_Output{
 			Properties: &_struct.Struct{
 				Fields: map[string]*_struct.Value{
@@ -53,10 +61,10 @@
 
 func TestCreateCombinedTestPlan_success(t *testing.T) {
 	reefGceTestCfg := &testplans.GceTestCfg{GceTest: []*testplans.GceTestCfg_GceTest{
-		{TestType: "GCE reef"},
+		{TestType: "GCE reef", Common: &testplans.TestSuiteCommon{Critical: &wrappers.BoolValue{Value: true}}},
 	}}
 	reefMoblabVmTestCfg := &testplans.MoblabVmTestCfg{MoblabTest: []*testplans.MoblabVmTestCfg_MoblabTest{
-		{TestType: "Moblab reef"},
+		{TestType: "Moblab reef", Common: &testplans.TestSuiteCommon{Critical: &wrappers.BoolValue{Value: true}}},
 	}}
 	kevinHWTestCfg := &testplans.HwTestCfg{HwTest: []*testplans.HwTestCfg_HwTest{
 		{
@@ -90,10 +98,10 @@
 	bbBuilds := []*bbproto.Build{
 		makeBuildbucketBuild("kevin", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
 			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-		}),
+		}, false),
 		makeBuildbucketBuild("reef", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
 			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-		}),
+		}, true),
 	}
 	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{
 		{
@@ -150,9 +158,53 @@
 					ArtifactsGsBucket: GS_BUCKET,
 					ArtifactsGsPath:   GS_PATH_PREFIX + "kevin",
 				}},
+		},
+		GceTestUnits: []*testplans.GceTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "reef",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "reef"}},
+				GceTestCfg: reefGceTestCfg},
+		},
+		MoblabVmTestUnits: []*testplans.MoblabVmTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "reef",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "reef"}},
+				MoblabVmTestCfg: reefMoblabVmTestCfg},
+		},
+		HwTestUnits: []*testplans.HwTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "kevin",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "kevin"}},
+				HwTestCfg: kevinHWTestCfg},
+		},
+		TastVmTestUnits: []*testplans.TastVmTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "kevin",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "kevin"}},
+				TastVmTestCfg: kevinTastVMTestCfg},
+		},
+		VmTestUnits: []*testplans.VmTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "kevin",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "kevin"}},
+				VmTestCfg: kevinVMTestCfg},
 		}}
-
-	if diff := cmp.Diff(expectedTestPlan, actualTestPlan); diff != "" {
+	if diff := cmp.Diff(expectedTestPlan, actualTestPlan, cmpopts.EquateEmpty()); diff != "" {
 		t.Errorf("CreateCombinedTestPlan bad result (-want/+got)\n%s", diff)
 	}
 }
@@ -184,10 +236,10 @@
 	bbBuilds := []*bbproto.Build{
 		makeBuildbucketBuild("kevin", bbproto.Status_FAILURE, []*bbproto.GerritChange{
 			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-		}),
+		}, true),
 		makeBuildbucketBuild("reef", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
 			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-		}),
+		}, true),
 	}
 	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{
 		{
@@ -216,9 +268,18 @@
 					ArtifactsGsBucket: GS_BUCKET,
 					ArtifactsGsPath:   GS_PATH_PREFIX + "reef",
 				}},
+		},
+		GceTestUnits: []*testplans.GceTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "reef",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "reef"}},
+				GceTestCfg: reefGceTestCfg},
 		}}
 
-	if diff := cmp.Diff(expectedTestPlan, actualTestPlan); diff != "" {
+	if diff := cmp.Diff(expectedTestPlan, actualTestPlan, cmpopts.EquateEmpty()); diff != "" {
 		t.Errorf("CreateCombinedTestPlan bad result (-want/+got)\n%s", diff)
 	}
 }
@@ -244,7 +305,7 @@
 	bbBuilds := []*bbproto.Build{
 		makeBuildbucketBuild("kevin", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
 			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-		}),
+		}, true),
 	}
 	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{
 		{
@@ -267,7 +328,7 @@
 	expectedTestPlan := &testplans.GenerateTestPlanResponse{
 		TestUnit: []*testplans.TestUnit{}}
 
-	if diff := cmp.Diff(expectedTestPlan, actualTestPlan); diff != "" {
+	if diff := cmp.Diff(expectedTestPlan, actualTestPlan, cmpopts.EquateEmpty()); diff != "" {
 		t.Errorf("CreateCombinedTestPlan bad result (-want/+got)\n%s", diff)
 	}
 }
@@ -308,7 +369,7 @@
 				TestRestriction: &testplans.TestRestriction{DisableHwTests: true}}}}
 	bbBuild := makeBuildbucketBuild("kevin", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
 		{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
-	})
+	}, true)
 	bbBuild.Output.Properties.Fields["pointless_build"] = &_struct.Value{Kind: &_struct.Value_BoolValue{BoolValue: true}}
 	bbBuilds := []*bbproto.Build{bbBuild}
 	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{
@@ -332,7 +393,7 @@
 	expectedTestPlan := &testplans.GenerateTestPlanResponse{
 		TestUnit: []*testplans.TestUnit{}}
 
-	if diff := cmp.Diff(expectedTestPlan, actualTestPlan); diff != "" {
+	if diff := cmp.Diff(expectedTestPlan, actualTestPlan, cmpopts.EquateEmpty()); diff != "" {
 		t.Errorf("CreateCombinedTestPlan bad result (-want/+got)\n%s", diff)
 	}
 }
@@ -342,7 +403,7 @@
 	sourceTreeTestCfg := &testplans.SourceTreeTestCfg{}
 	bbBuilds := []*bbproto.Build{
 		// build target is empty.
-		makeBuildbucketBuild("", bbproto.Status_FAILURE, []*bbproto.GerritChange{}),
+		makeBuildbucketBuild("", bbproto.Status_FAILURE, []*bbproto.GerritChange{}, true),
 	}
 	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{})
 	repoToSrcRoot := map[string]string{}