test_planner: implement new only-test functionality

This lets us say something like "if a CQ run only affects
chromite, then only run tests on kevin".

I had to refactor things a bit to express this sort-of cleanly.

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

Change-Id: I59aefb9ce082990d53d7f8b5adff1643e7c99bfd
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/infra/test_planner/+/1673802
Reviewed-by: Sean Abraham <seanabraham@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 4e2bef7..44799b9 100644
--- a/deploy_cipd.json
+++ b/deploy_cipd.json
@@ -1,6 +1,6 @@
 {
   "result": {
     "package": "chromiumos/infra/test_planner",
-    "instance_id": "rg7uZMjHiaUP2IrPUc5kyj_lH4XzP7fF7G4-AlEDImkC"
+    "instance_id": "9d0hPmn-M6kF0rTidWXJKuUa-w-mHLZI3W-kIzXl0bwC"
   }
 }
\ No newline at end of file
diff --git a/src/testplans/internal/generator/generator.go b/src/testplans/internal/generator/generator.go
index a111f6d..999f9d4 100644
--- a/src/testplans/internal/generator/generator.go
+++ b/src/testplans/internal/generator/generator.go
@@ -15,7 +15,6 @@
 	"testplans/internal/git"
 )
 
-
 // BuildTarget is an OS build target, such as "kevin" or "eve".
 type BuildTarget string
 
@@ -52,8 +51,15 @@
 		}
 	}
 
-	// BuildTargets for which HW or VM testing may be skipped, due to source tree configuration.
-	skippableTests, err := extractSkippableTests(sourceTreeCfg, filteredBbBuilds, changeRevs, repoToSrcRoot)
+	// Get the GerritChanges from any of the filtered builds, since the list
+	// should be the same for all of them.
+	changes := make([]*bbproto.GerritChange, 0)
+	if len(filteredBbBuilds) > 0 {
+		changes = append(changes, filteredBbBuilds[0].Input.GerritChanges...)
+	}
+
+	// For those changes, what pruning optimizations can be done?
+	pruneResult, err := extractPruneResult(sourceTreeCfg, changes, changeRevs, repoToSrcRoot)
 	if err != nil {
 		return testPlan, err
 	}
@@ -73,7 +79,7 @@
 		targetBuildResults = append(targetBuildResults, *tbr)
 	}
 
-	return createResponse(targetBuildResults, skippableTests)
+	return createResponse(targetBuildResults, pruneResult)
 }
 
 func isPointlessBuild(bb *bbproto.Build) bool {
@@ -125,7 +131,7 @@
 // createResponse creates the final GenerateTestPlanResponse.
 func createResponse(
 	targetBuildResults []targetBuildResult,
-	skippableTests map[BuildTarget]map[testType]bool) (*testplans.GenerateTestPlanResponse, error) {
+	pruneResult *testPruneResult) (*testplans.GenerateTestPlanResponse, error) {
 
 	resp := &testplans.GenerateTestPlanResponse{}
 targetLoop:
@@ -164,8 +170,10 @@
 				GceTestCfg: pttr.GceTestCfg})
 		}
 		if pttr.HwTestCfg != nil {
-			if skippableTests[tbr.buildTarget][hw] {
+			if pruneResult.disableHWTests {
 				log.Printf("No HW testing needed for %s", tbr.buildTarget)
+			} else if pruneResult.canSkipForOnlyTestRule(tbr.buildTarget) {
+				log.Printf("Using OnlyTest rule to skip HW testing for %s", tbr.buildTarget)
 			} else {
 				for _, hw := range pttr.HwTestCfg.HwTest {
 					hw.Common = withCritical(hw.Common, critical)
@@ -192,7 +200,7 @@
 				TastVmTestCfg: pttr.TastVmTestCfg})
 		}
 		if pttr.VmTestCfg != nil {
-			if skippableTests[tbr.buildTarget][vm] {
+			if pruneResult.disableVMTests {
 				log.Printf("No VM testing needed for %s", tbr.buildTarget)
 			} else {
 				for _, vm := range pttr.VmTestCfg.VmTest {
diff --git a/src/testplans/internal/generator/generator_test.go b/src/testplans/internal/generator/generator_test.go
index 0b404ea..b2e4e2b 100644
--- a/src/testplans/internal/generator/generator_test.go
+++ b/src/testplans/internal/generator/generator_test.go
@@ -336,6 +336,79 @@
 	}
 }
 
+func TestCreateCombinedTestPlan_doesOnlyTest(t *testing.T) {
+	kevinHWTestCfg := &testplans.HwTestCfg{HwTest: []*testplans.HwTestCfg_HwTest{
+		{
+			Suite:       "HW kevin",
+			SkylabBoard: "kev",
+		},
+	}}
+	bobHWTestCfg := &testplans.HwTestCfg{HwTest: []*testplans.HwTestCfg_HwTest{
+		{
+			Suite:       "HW bob",
+			SkylabBoard: "bob board",
+		},
+	}}
+	testReqs := &testplans.TargetTestRequirementsCfg{
+		PerTargetTestRequirements: []*testplans.PerTargetTestRequirements{
+			{TargetCriteria: &testplans.TargetCriteria{
+				TargetType: &testplans.TargetCriteria_BuildTarget{BuildTarget: "kevin"}},
+				HwTestCfg: kevinHWTestCfg},
+			{TargetCriteria: &testplans.TargetCriteria{
+				TargetType: &testplans.TargetCriteria_BuildTarget{BuildTarget: "bob"}},
+				HwTestCfg: bobHWTestCfg},
+		},
+	}
+	sourceTreeTestCfg := &testplans.SourceTreeTestCfg{
+		SourceTreeTestRestriction: []*testplans.SourceTreeTestRestriction{
+			{SourceTree: &testplans.SourceTree{Path: "no/hw/tests/here/some/file"},
+				TestRestriction: &testplans.TestRestriction{
+					OnlyTestBuildTargets: []*chromiumos.BuildTarget{
+						{Name: "kevin"}},
+				}}}}
+	bbBuilds := []*bbproto.Build{
+		makeBuildbucketBuild("kevin", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
+			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
+		}, true),
+		makeBuildbucketBuild("bob", bbproto.Status_SUCCESS, []*bbproto.GerritChange{
+			{Host: "test-review.googlesource.com", Change: 123, Patchset: 2},
+		}, true),
+	}
+	chRevData := git.GetChangeRevsForTest([]*git.ChangeRev{
+		{
+			ChangeRevKey: git.ChangeRevKey{
+				Host:      "test-review.googlesource.com",
+				ChangeNum: 123,
+				Revision:  2,
+			},
+			Project: "chromiumos/test/repo/name",
+			Files:   []string{"some/file"},
+		},
+	})
+	repoToSrcRoot := map[string]string{"chromiumos/test/repo/name": "no/hw/tests/here"}
+
+	actualTestPlan, err := CreateTestPlan(testReqs, sourceTreeTestCfg, bbBuilds, chRevData, repoToSrcRoot)
+	if err != nil {
+		t.Error(err)
+	}
+
+	expectedTestPlan := &testplans.GenerateTestPlanResponse{
+		HwTestUnits: []*testplans.HwTestUnit{
+			{Common: &testplans.TestUnitCommon{
+				BuildPayload: &testplans.BuildPayload{
+					ArtifactsGsBucket: GS_BUCKET,
+					ArtifactsGsPath:   GS_PATH_PREFIX + "kevin",
+				},
+				BuildTarget: &chromiumos.BuildTarget{Name: "kevin"}},
+				HwTestCfg: kevinHWTestCfg},
+		},
+	}
+
+	if diff := cmp.Diff(expectedTestPlan, actualTestPlan, cmpopts.EquateEmpty()); diff != "" {
+		t.Errorf("CreateCombinedTestPlan bad result (-want/+got)\n%s", diff)
+	}
+}
+
 func TestCreateCombinedTestPlan_inputMissingTargetType(t *testing.T) {
 	testReqs := &testplans.TargetTestRequirementsCfg{
 		PerTargetTestRequirements: []*testplans.PerTargetTestRequirements{
diff --git a/src/testplans/internal/generator/test_pruner.go b/src/testplans/internal/generator/test_pruner.go
index 0ec61bf..abf0bad 100644
--- a/src/testplans/internal/generator/test_pruner.go
+++ b/src/testplans/internal/generator/test_pruner.go
@@ -31,74 +31,134 @@
 	return [...]string{"Hw", "Vm"}[tt]
 }
 
-// extractSkippableTests maps BuildTargets to the test types that can be skipped for those targets,
-// based on source tree test restrictions.
-func extractSkippableTests(
-	sourceTreeCfg *testplans.SourceTreeTestCfg,
-	buildReports []*bbproto.Build,
-	changeRevs *git.ChangeRevData,
-	repoToSrcRoot map[string]string) (map[BuildTarget]map[testType]bool, error) {
-
-	skippableTests := make(map[BuildTarget]map[testType]bool)
-	for _, report := range buildReports {
-		buildTarget := BuildTarget(getBuildTarget(report))
-		skippableTests[buildTarget] = make(map[testType]bool)
-
-		disableHWTesting, err := canDisableTesting(sourceTreeCfg, report, changeRevs, repoToSrcRoot, hw)
-		if err != nil {
-			return skippableTests, err
-		}
-		skippableTests[buildTarget][hw] = disableHWTesting
-		disableVMTesting, err := canDisableTesting(sourceTreeCfg, report, changeRevs, repoToSrcRoot, vm)
-		if err != nil {
-			return skippableTests, err
-		}
-		skippableTests[buildTarget][vm] = disableVMTesting
-		log.Printf("For build %s, got disableHWTesting = %t, disableVMTesting = %t", getBuildTarget(report), disableHWTesting, disableVMTesting)
-	}
-	return skippableTests, nil
+type testPruneResult struct {
+	disableHWTests       bool
+	disableVMTests       bool
+	onlyTestBuildTargets map[BuildTarget]bool
 }
 
-
-// canDisableTesting determines whether a particular testing type is unnecessary for a given
-// builder, based on source tree test restrictions.
-func canDisableTesting(
-	sourceTreeCfg *testplans.SourceTreeTestCfg,
-	buildResult *bbproto.Build,
-	changeRevs *git.ChangeRevData,
-	repoToSrcRoot map[string]string,
-	tt testType) (bool, error) {
-	if len(buildResult.Input.GerritChanges) == 0 {
-		// This happens during postsubmit runs, for example.
-		log.Printf("build doesn't contain gerrit_changes %s, so no tests will be skipped", getBuildTarget(buildResult))
-		return false, nil
+// canSkipForOnlyTestRule identifies whether testing for a provided build target
+// can be skipped due to the only-test rules. e.g. if we only need to test on
+// "reef", this will return false for all non-reef build targets.
+func (tpr testPruneResult) canSkipForOnlyTestRule(bt BuildTarget) bool {
+	// If no only-test build targets were specified, we can't skip testing for
+	// any build targets by only-test rules.
+	if len(tpr.onlyTestBuildTargets) == 0 {
+		return false
 	}
-	fileCount := 0
-	for _, commit := range buildResult.Input.GerritChanges {
+	isAnOnlyTestTarget, _ := tpr.onlyTestBuildTargets[bt]
+	return !isAnOnlyTestTarget
+}
+
+func extractPruneResult(
+	sourceTreeCfg *testplans.SourceTreeTestCfg,
+	changes []*bbproto.GerritChange,
+	changeRevs *git.ChangeRevData,
+	repoToSrcRoot map[string]string) (*testPruneResult, error) {
+
+	result := &testPruneResult{}
+
+	if len(changes) == 0 {
+		// This happens during postsubmit runs, for example.
+		log.Print("no gerrit_changes, so no tests will be skipped")
+		return result, nil
+	}
+
+	// All the files in the GerritChanges, in source tree form.
+	srcPaths, err := srcPaths(changes, changeRevs, repoToSrcRoot)
+	if err != nil {
+		return result, err
+	}
+
+	disableHW := true
+	for _, fileSrcPath := range srcPaths {
+		if disableHW {
+			disableHWForPath, err := canDisableTestingForPath(fileSrcPath, sourceTreeCfg, hw)
+			if err != nil {
+				return result, err
+			}
+			if !disableHWForPath {
+				log.Printf("cannot disable HW testing due to file %s", fileSrcPath)
+				disableHW = false
+			}
+		}
+	}
+	disableVM := true
+	for _, fileSrcPath := range srcPaths {
+		if disableVM {
+			disableVMForPath, err := canDisableTestingForPath(fileSrcPath, sourceTreeCfg, vm)
+			if err != nil {
+				return result, err
+			}
+			if !disableVMForPath {
+				log.Printf("cannot disable VM testing due to file %s", fileSrcPath)
+				disableVM = false
+			}
+		}
+	}
+	canOnlyTestSomeBuildTargets := true
+	onlyTestBuildTargets := make(map[BuildTarget]bool)
+	for _, fileSrcPath := range srcPaths {
+		if canOnlyTestSomeBuildTargets {
+			fileOnlyTestBuildTargets, err := checkOnlyTestBuildTargets(fileSrcPath, sourceTreeCfg)
+			if err != nil {
+				return result, err
+			}
+			if len(fileOnlyTestBuildTargets) == 0 {
+				log.Printf("cannot limit set of build targets for testing due to %s", fileSrcPath)
+				canOnlyTestSomeBuildTargets = false
+				onlyTestBuildTargets = make(map[BuildTarget]bool)
+			} else {
+				for k, v := range fileOnlyTestBuildTargets {
+					onlyTestBuildTargets[k] = v
+				}
+			}
+		}
+	}
+	return &testPruneResult{
+			disableHWTests:       disableHW,
+			disableVMTests:       disableVM,
+			onlyTestBuildTargets: onlyTestBuildTargets},
+		nil
+}
+
+// srcPaths extracts the source paths from each of the provided Gerrit changes.
+func srcPaths(
+	changes []*bbproto.GerritChange,
+	changeRevs *git.ChangeRevData,
+	repoToSrcRoot map[string]string) ([]string, error) {
+	srcPaths := make([]string, 0)
+	for _, commit := range changes {
 		chRev, err := changeRevs.GetChangeRev(commit.Host, commit.Change, int32(commit.Patchset))
 		if err != nil {
-			return false, err
+			return srcPaths, err
 		}
 		for _, file := range chRev.Files {
-			fileCount++
 			srcRootMapping, found := repoToSrcRoot[chRev.Project]
 			if !found {
-				return false, fmt.Errorf("Found no source mapping for project %s", chRev.Project)
+				return srcPaths, fmt.Errorf("Found no source mapping for project %s", chRev.Project)
 			}
-			fileSrcPath := fmt.Sprintf("%s/%s", srcRootMapping, file)
-			disableTesting, err := canDisableTestingForPath(fileSrcPath, sourceTreeCfg, tt)
-			if err != nil {
-				return false, err
-			}
-			if !disableTesting {
-				log.Printf("Can't disable %s testing due to file %s", tt, fileSrcPath)
-				return false, nil
+			srcPaths = append(srcPaths, fmt.Sprintf("%s/%s", srcRootMapping, file))
+		}
+	}
+	return srcPaths, nil
+}
+
+// checkOnlyTestBuildTargets checks if the provided path is covered by an
+// only-test rule, which would allow us to exclude testing for all other
+// build targets.
+func checkOnlyTestBuildTargets(
+	sourceTreePath string,
+	sourceTreeCfg *testplans.SourceTreeTestCfg) (map[BuildTarget]bool, error) {
+	result := make(map[BuildTarget]bool)
+	for _, sourceTreeRestriction := range sourceTreeCfg.SourceTreeTestRestriction {
+		if hasPathPrefix(sourceTreePath, sourceTreeRestriction.SourceTree.Path) {
+			for _, otbt := range sourceTreeRestriction.TestRestriction.OnlyTestBuildTargets {
+				result[BuildTarget(otbt.Name)] = true
 			}
 		}
 	}
-	// Either testing is disabled for all of the files or there are zero files.
-	log.Printf("%s testing is not needed for %d files in the %s build", tt, fileCount, getBuildTarget(buildResult))
-	return true, nil
+	return result, nil
 }
 
 // canDisableTestingForPath determines whether a particular testing type is unnecessary for
@@ -110,10 +170,21 @@
 			return false, fmt.Errorf("Missing test filter for %v", tt)
 		}
 		if testFilter(sourceTreeRestriction.TestRestriction) {
-			if strings.HasPrefix(sourceTreePath, sourceTreeRestriction.SourceTree.Path) {
+			if hasPathPrefix(sourceTreePath, sourceTreeRestriction.SourceTree.Path) {
 				return true, nil
 			}
 		}
 	}
 	return false, nil
 }
+
+// hasPathPrefix checks if the provided string has a provided path prefix.
+// e.g. ab/cd/ef, ab --> true
+//      ab/cd, ab/c --> false
+func hasPathPrefix(s string, prefix string) bool {
+	if s == prefix {
+		return true
+	}
+	prefixAsDir := strings.TrimSuffix(prefix, "/") + "/"
+	return strings.HasPrefix(s, prefixAsDir)
+}