[testsharder] Extend testsharder to take a test multipliers json file.

This file should contain tests that are to be run multiple times which the testsharder
will use to create a separate shard per test that contains the same test written as many
times as the total_runs specified in the file.

The output will contain the original shards it currently produces plus
these new shards with the multiplied tests. This will be used to test for flakiness.

Bug: IN-1231 #comment
Change-Id: I4366ff012bc929c938c6aeaa9a26812e7888c7ef
diff --git a/cmd/testsharder/main.go b/cmd/testsharder/main.go
index 5518192..54aebd8 100644
--- a/cmd/testsharder/main.go
+++ b/cmd/testsharder/main.go
@@ -26,6 +26,9 @@
 
 	// Tags are keys on which to filter environments, which are labeled.
 	tags command.StringsFlag
+
+	// The path to the json manifest file containing the tests to mutiply.
+	multipliersPath string
 )
 
 func usage() {
@@ -42,6 +45,7 @@
 	flag.StringVar(&outputFile, "output-file", "", "path to a file which will contain the shards as JSON, default is stdout")
 	flag.Var(&mode, "mode", "mode in which to run the testsharder (e.g., normal or restricted).")
 	flag.Var(&tags, "tag", "environment tags on which to filter; only the tests that match all tags will be sharded")
+	flag.StringVar(&multipliersPath, "multipliers", "", "path to the json manifest containing tests to multiply")
 	flag.Usage = usage
 }
 
@@ -68,6 +72,13 @@
 
 	// Create shards and write them to an output file if specifed, else stdout.
 	shards := testsharder.MakeShards(specs, mode, tags)
+	if multipliersPath != "" {
+		multipliers, err := testsharder.LoadTestModifiers(multipliersPath)
+		if err != nil {
+			log.Fatal(err)
+		}
+		shards = testsharder.MultiplyShards(shards, multipliers)
+	}
 	f := os.Stdout
 	if outputFile != "" {
 		var err error
diff --git a/testsharder/shard.go b/testsharder/shard.go
index a0a598e..8b09a3e 100644
--- a/testsharder/shard.go
+++ b/testsharder/shard.go
@@ -72,6 +72,34 @@
 	return shards
 }
 
+// Appends new shards to shards where each new shard contains one test repeated
+// multiple times according to the specifications in multipliers.
+func MultiplyShards(shards []*Shard, multipliers []TestModifier) []*Shard {
+	for _, shard := range shards {
+		for _, multiplier := range multipliers {
+			for _, test := range shard.Tests {
+				if multiplier.Target == test.Name && multiplier.OS == test.OS {
+					shards = append(shards, &Shard{
+						Name:  shard.Env.Name() + " - " + test.Name,
+						Tests: multiplyTest(test, multiplier.TotalRuns),
+						Env:   shard.Env,
+					})
+				}
+			}
+		}
+	}
+	return shards
+}
+
+// Returns a list of Tests containing the same test multiplied by the number of runs.
+func multiplyTest(test Test, runs int) []Test {
+	var tests []Test
+	for i := 0; i < runs; i++ {
+		tests = append(tests, test)
+	}
+	return tests
+}
+
 // Abstracts a mapping Environment -> []string, as Environment contains non-comparable
 // members (e.g., string slices), which makes it invalid for a map key.
 type envMap struct {
diff --git a/testsharder/shard_test.go b/testsharder/shard_test.go
index 6a00c73..c34e702 100644
--- a/testsharder/shard_test.go
+++ b/testsharder/shard_test.go
@@ -172,3 +172,81 @@
 		assertEqual(t, expected, actual)
 	})
 }
+
+func TestMultiplyShards(t *testing.T) {
+	env1 := Environment{
+		Dimensions: DimensionSet{DeviceType: "QEMU"},
+		Tags:       []string{},
+	}
+	env2 := Environment{
+		Dimensions: DimensionSet{DeviceType: "NUC"},
+		Tags:       []string{},
+	}
+	env3 := Environment{
+		Dimensions: DimensionSet{OS: "Linux"},
+		Tags:       []string{},
+	}
+	makeTest := func(id int, os OS) Test {
+		return Test{
+			Name:     fmt.Sprintf("test%d", id),
+			Location: fmt.Sprintf("/path/to/test/%d", id),
+			OS:       os,
+		}
+	}
+
+	shard := func(env Environment, os OS, ids ...int) *Shard {
+		var tests []Test
+		for _, id := range ids {
+			tests = append(tests, makeTest(id, os))
+		}
+		return &Shard{
+			Name:  env.Name(),
+			Tests: tests,
+			Env:   env,
+		}
+	}
+
+	makeTestModifier := func(id int, os OS, runs int) TestModifier {
+		return TestModifier{
+			Target:    fmt.Sprintf("test%d", id),
+			OS:        os,
+			TotalRuns: runs,
+		}
+	}
+
+	multShard := func(env Environment, os OS, id int, runs int) *Shard {
+		var tests []Test
+		test := makeTest(id, os)
+		for i := 0; i < runs; i++ {
+			tests = append(tests, test)
+		}
+		return &Shard{
+			Name:  env.Name() + " - " + test.Name,
+			Tests: tests,
+			Env:   env,
+		}
+	}
+
+	t.Run("multiply tests in shards", func(t *testing.T) {
+		shards := []*Shard{
+			shard(env1, Fuchsia, 1),
+			shard(env2, Fuchsia, 1, 2),
+			shard(env3, Linux, 3),
+		}
+		multipliers := []TestModifier{
+			makeTestModifier(1, Fuchsia, 2),
+			makeTestModifier(3, Linux, 3),
+		}
+		actual := MultiplyShards(
+			shards,
+			multipliers,
+		)
+		expected := append(
+			shards,
+			multShard(env1, Fuchsia, 1, 2),
+			multShard(env2, Fuchsia, 1, 2),
+			multShard(env3, Linux, 3, 3),
+		)
+		assertEqual(t, expected, actual)
+	})
+}
diff --git a/testsharder/test_modifier.go b/testsharder/test_modifier.go
new file mode 100644
index 0000000..fa3b9c0
--- /dev/null
+++ b/testsharder/test_modifier.go
@@ -0,0 +1,48 @@
+// Copyright 2019 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+package testsharder
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+)
+
+// TestModifier is the specification for a single test and the number of
+// times it should be run.
+type TestModifier struct {
+	// Target is the GN target name of the test.
+	Target string `json:"target"`
+
+	// OS is the operating system in which this test must be executed; treated as "fuchsia" if not present.
+	OS OS `json:"os,omitempty"`
+
+	// TotalRuns is the number of times to run the test; treated as 1 if not present.
+	TotalRuns int `json:"total_runs,omitempty"`
+}
+
+// LoadTestModifiers loads a set of test modifiers from a json manifest.
+func LoadTestModifiers(manifestPath string) ([]TestModifier, error) {
+	bytes, err := ioutil.ReadFile(manifestPath)
+	if err != nil {
+		return nil, err
+	}
+	var specs []TestModifier
+	if err = json.Unmarshal(bytes, &specs); err != nil {
+		return nil, err
+	}
+
+	for i, _ := range specs {
+		if specs[i].Target == "" {
+			return nil, fmt.Errorf("A test spec's target must have a non-empty name")
+		}
+		if specs[i].TotalRuns == 0 {
+			specs[i].TotalRuns = 1
+		}
+		if specs[i].OS == "" {
+			specs[i].OS = Fuchsia
+		}
+	}
+	return specs, nil
+}
diff --git a/testsharder/test_modifier_test.go b/testsharder/test_modifier_test.go
new file mode 100644
index 0000000..59785b1
--- /dev/null
+++ b/testsharder/test_modifier_test.go
@@ -0,0 +1,76 @@
+// Copyright 2019 The Fuchsia Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+package testsharder
+
+import (
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"reflect"
+	"sort"
+	"testing"
+)
+
+var barTestModifier = TestModifier{
+	Target:    "//obsidian/lib/bar:bar_tests",
+	TotalRuns: 2,
+}
+
+var bazTestModifier = TestModifier{
+	Target: "//obsidian/public/lib/baz:baz_host_tests",
+	OS:     Linux,
+}
+
+func TestLoadTestModifiers(t *testing.T) {
+	areEqual := func(a, b []TestModifier) bool {
+		stringify := func(modifier TestModifier) string {
+			return fmt.Sprintf("%#v", modifier)
+		}
+		sort := func(list []TestModifier) {
+			sort.Slice(list[:], func(i, j int) bool {
+				return stringify(list[i]) < stringify(list[j])
+			})
+		}
+		sort(a)
+		sort(b)
+		return reflect.DeepEqual(a, b)
+	}
+
+	tmpDir, err := ioutil.TempDir("", "test-spec")
+	if err != nil {
+		t.Fatalf("failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	initial := []TestModifier{barTestModifier, bazTestModifier}
+
+	modifiersPath := filepath.Join(tmpDir, "test_modifiers.json")
+	m, err := os.Create(modifiersPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer m.Close()
+	if err := json.NewEncoder(m).Encode(&initial); err != nil {
+		t.Fatal(err)
+	}
+
+	actual, err := LoadTestModifiers(modifiersPath)
+	if err != nil {
+		t.Fatalf("failed to load test modifiers: %v", err)
+	}
+
+	bazOut := bazTestModifier
+	// If TotalRuns is missing, it gets set to default 1.
+	bazOut.TotalRuns = 1
+	barOut := barTestModifier
+	// If OS is missing, it gets set to default Fuchsia.
+	barOut.OS = Fuchsia
+	expected := []TestModifier{barOut, bazOut}
+
+	if !areEqual(expected, actual) {
+		t.Fatalf("test modifiers not properly loaded:\nexpected:\n%+v\nactual:\n%+v", expected, actual)
+	}
+}