Allow adding extra tradefed options in the Android.bp file

Some tests need to add custom tradefed options, but still want to
keep most of the soong autogenerated tradefed xml file.

Expose a test_options: { tradefed_options: [...] } property that
will allow tests to add more options to the autogenerated xml file.

Fixes: 184895128
Test: go test, and verified that the ninja files did not change for aosp_arm64
Change-Id: I75f7eb002c8325ce7cdc76e12e76e16195320620

10 files changed, 221 insertions, 166 deletions

diff --git a/android/module.go b/android/module.go
index 681f724b0..bf9737aa2 100644
--- a/android/module.go
+++ b/android/module.go
@@ -27,7 +27,6 @@ import (
 	"text/scanner"
 
 	"android/soong/bazel"
-
 	"github.com/google/blueprint"
 	"github.com/google/blueprint/proptools"
 )
diff --git a/cc/test.go b/cc/test.go
index 536210ba4..ed2a58a04 100644
--- a/cc/test.go
+++ b/cc/test.go
@@ -459,8 +459,16 @@ func (test *testBinary) install(ctx ModuleContext, file android.Path) {
 		configs = append(configs, tradefed.Object{"module_controller", "com.android.tradefed.testtype.suite.module.MinApiLevelModuleController", options})
 	}
 
-	test.testConfig = tradefed.AutoGenNativeTestConfig(ctx, test.Properties.Test_config,
-		test.Properties.Test_config_template, test.testDecorator.InstallerProperties.Test_suites, configs, test.Properties.Auto_gen_config, testInstallBase)
+	test.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(test.Properties.Test_config).
+		SetTestTemplateConfigProp(test.Properties.Test_config_template).
+		SetTestSuites(test.testDecorator.InstallerProperties.Test_suites).
+		SetConfig(configs).
+		SetAutoGenConfig(test.Properties.Auto_gen_config).
+		SetTestInstallBase(testInstallBase).
+		SetDeviceTemplate("${NativeTestConfigTemplate}").
+		SetHostTemplate("${NativeHostTestConfigTemplate}").
+		Build()
 
 	test.extraTestConfigs = android.PathsForModuleSrc(ctx, test.Properties.Test_options.Extra_test_configs)
 
@@ -616,8 +624,15 @@ func (benchmark *benchmarkDecorator) install(ctx ModuleContext, file android.Pat
 	if Bool(benchmark.Properties.Require_root) {
 		configs = append(configs, tradefed.Object{"target_preparer", "com.android.tradefed.targetprep.RootTargetPreparer", nil})
 	}
-	benchmark.testConfig = tradefed.AutoGenNativeBenchmarkTestConfig(ctx, benchmark.Properties.Test_config,
-		benchmark.Properties.Test_config_template, benchmark.Properties.Test_suites, configs, benchmark.Properties.Auto_gen_config)
+	benchmark.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(benchmark.Properties.Test_config).
+		SetTestTemplateConfigProp(benchmark.Properties.Test_config_template).
+		SetTestSuites(benchmark.Properties.Test_suites).
+		SetConfig(configs).
+		SetAutoGenConfig(benchmark.Properties.Auto_gen_config).
+		SetDeviceTemplate("${NativeBenchmarkTestConfigTemplate}").
+		SetHostTemplate("${NativeBenchmarkTestConfigTemplate}").
+		Build()
 
 	benchmark.binaryDecorator.baseInstaller.dir = filepath.Join("benchmarktest", ctx.ModuleName())
 	benchmark.binaryDecorator.baseInstaller.dir64 = filepath.Join("benchmarktest64", ctx.ModuleName())
diff --git a/java/java.go b/java/java.go
index 9dd585062..dd2437621 100644
--- a/java/java.go
+++ b/java/java.go
@@ -888,6 +888,10 @@ type TestOptions struct {
 
 	// a list of extra test configuration files that should be installed with the module.
 	Extra_test_configs []string `android:"path,arch_variant"`
+
+	// Extra <option> tags to add to the auto generated test xml file. The "key"
+	// is optional in each of these.
+	Tradefed_options []tradefed.Option
 }
 
 type testProperties struct {
@@ -1166,8 +1170,18 @@ func (j *Test) generateAndroidBuildActionsWithConfig(ctx android.ModuleContext,
 		j.testProperties.Test_options.Unit_test = proptools.BoolPtr(defaultUnitTest)
 	}
 
-	j.testConfig = tradefed.AutoGenJavaTestConfig(ctx, j.testProperties.Test_config, j.testProperties.Test_config_template,
-		j.testProperties.Test_suites, configs, j.testProperties.Auto_gen_config, j.testProperties.Test_options.Unit_test)
+	j.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(j.testProperties.Test_config).
+		SetTestTemplateConfigProp(j.testProperties.Test_config_template).
+		SetTestSuites(j.testProperties.Test_suites).
+		SetConfig(configs).
+		SetOptionsForAutogenerated(j.testProperties.Test_options.Tradefed_options).
+		SetAutoGenConfig(j.testProperties.Auto_gen_config).
+		SetUnitTest(j.testProperties.Test_options.Unit_test).
+		SetDeviceTemplate("${JavaTestConfigTemplate}").
+		SetHostTemplate("${JavaHostTestConfigTemplate}").
+		SetHostUnitTestTemplate("${JavaHostUnitTestConfigTemplate}").
+		Build()
 
 	j.data = android.PathsForModuleSrc(ctx, j.testProperties.Data)
 
@@ -1212,8 +1226,13 @@ func (j *TestHelperLibrary) GenerateAndroidBuildActions(ctx android.ModuleContex
 }
 
 func (j *JavaTestImport) GenerateAndroidBuildActions(ctx android.ModuleContext) {
-	j.testConfig = tradefed.AutoGenJavaTestConfig(ctx, j.prebuiltTestProperties.Test_config, nil,
-		j.prebuiltTestProperties.Test_suites, nil, nil, nil)
+	j.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(j.prebuiltTestProperties.Test_config).
+		SetTestSuites(j.prebuiltTestProperties.Test_suites).
+		SetDeviceTemplate("${JavaTestConfigTemplate}").
+		SetHostTemplate("${JavaHostTestConfigTemplate}").
+		SetHostUnitTestTemplate("${JavaHostUnitTestConfigTemplate}").
+		Build()
 
 	j.Import.GenerateAndroidBuildActions(ctx)
 }
diff --git a/java/java_test.go b/java/java_test.go
index dff1fd07f..62a372cd9 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -1945,3 +1945,25 @@ func TestJavaApiLibraryJarGeneration(t *testing.T) {
 		}
 	}
 }
+
+func TestTradefedOptions(t *testing.T) {
+	result := PrepareForTestWithJavaBuildComponents.RunTestWithBp(t, `
+java_test_host {
+	name: "foo",
+	test_options: {
+		tradefed_options: [
+			{
+				name: "exclude-path",
+				value: "org/apache"
+			}
+		]
+	}
+}
+`)
+	args := result.ModuleForTests("foo", "linux_glibc_common").
+		Output("out/soong/.intermediates/foo/linux_glibc_common/foo.config").Args
+	expected := proptools.NinjaAndShellEscape("<option name=\"exclude-path\" value=\"org/apache\" />")
+	if args["extraConfigs"] != expected {
+		t.Errorf("Expected args[\"extraConfigs\"] to equal %q, was %q", expected, args["extraConfigs"])
+	}
+}
diff --git a/java/robolectric.go b/java/robolectric.go
index 938abe138..6a2d0b3ac 100644
--- a/java/robolectric.go
+++ b/java/robolectric.go
@@ -131,9 +131,14 @@ func (r *robolectricTest) GenerateAndroidBuildActions(ctx android.ModuleContext)
 	r.forceOSType = ctx.Config().BuildOS
 	r.forceArchType = ctx.Config().BuildArch
 
-	r.testConfig = tradefed.AutoGenRobolectricTestConfig(ctx, r.testProperties.Test_config,
-		r.testProperties.Test_config_template, r.testProperties.Test_suites,
-		r.testProperties.Auto_gen_config)
+	r.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(r.testProperties.Test_config).
+		SetTestTemplateConfigProp(r.testProperties.Test_config_template).
+		SetTestSuites(r.testProperties.Test_suites).
+		SetAutoGenConfig(r.testProperties.Auto_gen_config).
+		SetDeviceTemplate("${RobolectricTestConfigTemplate}").
+		SetHostTemplate("${RobolectricTestConfigTemplate}").
+		Build()
 	r.data = android.PathsForModuleSrc(ctx, r.testProperties.Data)
 
 	roboTestConfig := android.PathForModuleGen(ctx, "robolectric").
diff --git a/python/test.go b/python/test.go
index b9b346549..5781df7c7 100644
--- a/python/test.go
+++ b/python/test.go
@@ -67,9 +67,14 @@ func (test *testDecorator) bootstrapperProps() []interface{} {
 }
 
 func (test *testDecorator) install(ctx android.ModuleContext, file android.Path) {
-	test.testConfig = tradefed.AutoGenPythonBinaryHostTestConfig(ctx, test.testProperties.Test_config,
-		test.testProperties.Test_config_template, test.binaryDecorator.binaryProperties.Test_suites,
-		test.binaryDecorator.binaryProperties.Auto_gen_config)
+	test.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(test.testProperties.Test_config).
+		SetTestTemplateConfigProp(test.testProperties.Test_config_template).
+		SetTestSuites(test.binaryDecorator.binaryProperties.Test_suites).
+		SetAutoGenConfig(test.binaryDecorator.binaryProperties.Auto_gen_config).
+		SetDeviceTemplate("${PythonBinaryHostTestConfigTemplate}").
+		SetHostTemplate("${PythonBinaryHostTestConfigTemplate}").
+		Build()
 
 	test.binaryDecorator.pythonInstaller.dir = "nativetest"
 	test.binaryDecorator.pythonInstaller.dir64 = "nativetest64"
diff --git a/rust/benchmark.go b/rust/benchmark.go
index 0e842435d..b417a2d65 100644
--- a/rust/benchmark.go
+++ b/rust/benchmark.go
@@ -112,12 +112,14 @@ func (benchmark *benchmarkDecorator) compilerProps() []interface{} {
 }
 
 func (benchmark *benchmarkDecorator) install(ctx ModuleContext) {
-	benchmark.testConfig = tradefed.AutoGenRustBenchmarkConfig(ctx,
-		benchmark.Properties.Test_config,
-		benchmark.Properties.Test_config_template,
-		benchmark.Properties.Test_suites,
-		nil,
-		benchmark.Properties.Auto_gen_config)
+	benchmark.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(benchmark.Properties.Test_config).
+		SetTestTemplateConfigProp(benchmark.Properties.Test_config_template).
+		SetTestSuites(benchmark.Properties.Test_suites).
+		SetAutoGenConfig(benchmark.Properties.Auto_gen_config).
+		SetDeviceTemplate("${RustDeviceBenchmarkConfigTemplate}").
+		SetHostTemplate("${RustHostBenchmarkConfigTemplate}").
+		Build()
 
 	// default relative install path is module name
 	if !Bool(benchmark.Properties.No_named_install_directory) {
diff --git a/rust/test.go b/rust/test.go
index 0cc3bca6e..ecc7d5d24 100644
--- a/rust/test.go
+++ b/rust/test.go
@@ -130,13 +130,16 @@ func (test *testDecorator) install(ctx ModuleContext) {
 		configs = append(configs, tradefed.Object{"target_preparer", "com.android.tradefed.targetprep.RootTargetPreparer", options})
 	}
 
-	test.testConfig = tradefed.AutoGenRustTestConfig(ctx,
-		test.Properties.Test_config,
-		test.Properties.Test_config_template,
-		test.Properties.Test_suites,
-		configs,
-		test.Properties.Auto_gen_config,
-		testInstallBase)
+	test.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(test.Properties.Test_config).
+		SetTestTemplateConfigProp(test.Properties.Test_config_template).
+		SetTestSuites(test.Properties.Test_suites).
+		SetConfig(configs).
+		SetAutoGenConfig(test.Properties.Auto_gen_config).
+		SetTestInstallBase(testInstallBase).
+		SetDeviceTemplate("${RustDeviceTestConfigTemplate}").
+		SetHostTemplate("${RustHostTestConfigTemplate}").
+		Build()
 
 	dataSrcPaths := android.PathsForModuleSrc(ctx, test.Properties.Data)
 
diff --git a/sh/sh_binary.go b/sh/sh_binary.go
index 96273297a..4eae397a6 100644
--- a/sh/sh_binary.go
+++ b/sh/sh_binary.go
@@ -379,8 +379,16 @@ func (s *ShTest) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 		}
 		configs = append(configs, tradefed.Object{"target_preparer", "com.android.tradefed.targetprep.PushFilePreparer", options})
 	}
-	s.testConfig = tradefed.AutoGenShellTestConfig(ctx, s.testProperties.Test_config,
-		s.testProperties.Test_config_template, s.testProperties.Test_suites, configs, s.testProperties.Auto_gen_config, s.outputFilePath.Base())
+	s.testConfig = tradefed.NewMaybeAutoGenTestConfigBuilder(ctx).
+		SetTestConfigProp(s.testProperties.Test_config).
+		SetTestTemplateConfigProp(s.testProperties.Test_config_template).
+		SetTestSuites(s.testProperties.Test_suites).
+		SetConfig(configs).
+		SetAutoGenConfig(s.testProperties.Auto_gen_config).
+		SetOutputFileName(s.outputFilePath.Base()).
+		SetDeviceTemplate("${ShellTestConfigTemplate}").
+		SetHostTemplate("${ShellTestConfigTemplate}").
+		Build()
 
 	s.dataModules = make(map[string]android.Path)
 	ctx.VisitDirectDeps(func(dep android.Module) {
diff --git a/tradefed/autogen.go b/tradefed/autogen.go
index c2429ab7a..236e559b5 100644
--- a/tradefed/autogen.go
+++ b/tradefed/autogen.go
@@ -107,176 +107,153 @@ func (ob Object) Config() string {
 
 }
 
-func autogenTemplate(ctx android.ModuleContext, output android.WritablePath, template string, configs []Config, testInstallBase string) {
-	autogenTemplateWithNameAndOutputFile(ctx, ctx.ModuleName(), output, template, configs, "", testInstallBase)
+// MaybeAutoGenTestConfigBuilder provides a Build() method that will either
+// generate a AndroidTest.xml file, or use an existing user-supplied one.
+// It used to be a bunch of separate functions for each language, but was
+// converted to this builder pattern to have one function that accepts many
+// optional arguments.
+type MaybeAutoGenTestConfigBuilder struct {
+	ctx                     android.ModuleContext
+	name                    string
+	outputFileName          string
+	testConfigProp          *string
+	testConfigTemplateProp  *string
+	testSuites              []string
+	config                  []Config
+	configsForAutogenerated []Config
+	autoGenConfig           *bool
+	unitTest                *bool
+	testInstallBase         string
+	deviceTemplate          string
+	hostTemplate            string
+	hostUnitTestTemplate    string
 }
 
-func autogenTemplateWithName(ctx android.ModuleContext, name string, output android.WritablePath, template string, configs []Config, testInstallBase string) {
-	autogenTemplateWithNameAndOutputFile(ctx, name, output, template, configs, "", testInstallBase)
+func NewMaybeAutoGenTestConfigBuilder(ctx android.ModuleContext) *MaybeAutoGenTestConfigBuilder {
+	return &MaybeAutoGenTestConfigBuilder{
+		ctx:  ctx,
+		name: ctx.ModuleName(),
+	}
 }
 
-func autogenTemplateWithNameAndOutputFile(ctx android.ModuleContext, name string, output android.WritablePath, template string, configs []Config, outputFileName string, testInstallBase string) {
-	var configStrings []string
-	for _, config := range configs {
-		configStrings = append(configStrings, config.Config())
-	}
-	extraConfigs := strings.Join(configStrings, fmt.Sprintf("\\n%s", test_xml_indent))
-	extraConfigs = proptools.NinjaAndShellEscape(extraConfigs)
+func (b *MaybeAutoGenTestConfigBuilder) SetName(name string) *MaybeAutoGenTestConfigBuilder {
+	b.name = name
+	return b
+}
 
-	ctx.Build(pctx, android.BuildParams{
-		Rule:        autogenTestConfig,
-		Description: "test config",
-		Output:      output,
-		Args: map[string]string{
-			"name":            name,
-			"template":        template,
-			"extraConfigs":    extraConfigs,
-			"outputFileName":  outputFileName,
-			"testInstallBase": testInstallBase,
-		},
-	})
+func (b *MaybeAutoGenTestConfigBuilder) SetOutputFileName(outputFileName string) *MaybeAutoGenTestConfigBuilder {
+	b.outputFileName = outputFileName
+	return b
 }
 
-func AutoGenNativeTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool, testInstallBase string) android.Path {
+func (b *MaybeAutoGenTestConfigBuilder) SetTestConfigProp(testConfigProp *string) *MaybeAutoGenTestConfigBuilder {
+	b.testConfigProp = testConfigProp
+	return b
+}
 
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), config, testInstallBase)
-		} else {
-			if ctx.Device() {
-				autogenTemplate(ctx, autogenPath, "${NativeTestConfigTemplate}", config, testInstallBase)
-			} else {
-				autogenTemplate(ctx, autogenPath, "${NativeHostTestConfigTemplate}", config, testInstallBase)
-			}
-		}
-		return autogenPath
-	}
-	return path
+func (b *MaybeAutoGenTestConfigBuilder) SetTestTemplateConfigProp(testConfigTemplateProp *string) *MaybeAutoGenTestConfigBuilder {
+	b.testConfigTemplateProp = testConfigTemplateProp
+	return b
 }
 
-func AutoGenShellTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool, outputFileName string) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplateWithNameAndOutputFile(ctx, ctx.ModuleName(), autogenPath, templatePath.String(), config, outputFileName, "")
-		} else {
-			autogenTemplateWithNameAndOutputFile(ctx, ctx.ModuleName(), autogenPath, "${ShellTestConfigTemplate}", config, outputFileName, "")
-		}
-		return autogenPath
-	}
-	return path
+func (b *MaybeAutoGenTestConfigBuilder) SetTestSuites(testSuites []string) *MaybeAutoGenTestConfigBuilder {
+	b.testSuites = testSuites
+	return b
 }
 
-func AutoGenNativeBenchmarkTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, configs []Config, autoGenConfig *bool) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), configs, "")
-		} else {
-			autogenTemplate(ctx, autogenPath, "${NativeBenchmarkTestConfigTemplate}", configs, "")
-		}
-		return autogenPath
-	}
-	return path
+func (b *MaybeAutoGenTestConfigBuilder) SetConfig(config []Config) *MaybeAutoGenTestConfigBuilder {
+	b.config = config
+	return b
 }
 
-func AutoGenJavaTestConfig(ctx android.ModuleContext, testConfigProp *string, testConfigTemplateProp *string,
-	testSuites []string, config []Config, autoGenConfig *bool, unitTest *bool) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), config, "")
-		} else {
-			if ctx.Device() {
-				autogenTemplate(ctx, autogenPath, "${JavaTestConfigTemplate}", config, "")
-			} else {
-				if Bool(unitTest) {
-					autogenTemplate(ctx, autogenPath, "${JavaHostUnitTestConfigTemplate}", config, "")
-				} else {
-					autogenTemplate(ctx, autogenPath, "${JavaHostTestConfigTemplate}", config, "")
-				}
-			}
-		}
-		return autogenPath
+func (b *MaybeAutoGenTestConfigBuilder) SetOptionsForAutogenerated(configsForAutogenerated []Option) *MaybeAutoGenTestConfigBuilder {
+	configs := make([]Config, 0, len(configsForAutogenerated))
+	for _, c := range configsForAutogenerated {
+		configs = append(configs, c)
 	}
-	return path
+	b.configsForAutogenerated = configs
+	return b
 }
 
-func AutoGenPythonBinaryHostTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, autoGenConfig *bool) android.Path {
+func (b *MaybeAutoGenTestConfigBuilder) SetUnitTest(unitTest *bool) *MaybeAutoGenTestConfigBuilder {
+	b.unitTest = unitTest
+	return b
+}
 
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), nil, "")
-		} else {
-			autogenTemplate(ctx, autogenPath, "${PythonBinaryHostTestConfigTemplate}", nil, "")
-		}
-		return autogenPath
-	}
-	return path
+func (b *MaybeAutoGenTestConfigBuilder) SetAutoGenConfig(autoGenConfig *bool) *MaybeAutoGenTestConfigBuilder {
+	b.autoGenConfig = autoGenConfig
+	return b
 }
 
-func AutoGenRustTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool, testInstallBase string) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), config, testInstallBase)
-		} else {
-			if ctx.Device() {
-				autogenTemplate(ctx, autogenPath, "${RustDeviceTestConfigTemplate}", config, testInstallBase)
-			} else {
-				autogenTemplate(ctx, autogenPath, "${RustHostTestConfigTemplate}", config, testInstallBase)
-			}
-		}
-		return autogenPath
-	}
-	return path
+func (b *MaybeAutoGenTestConfigBuilder) SetTestInstallBase(testInstallBase string) *MaybeAutoGenTestConfigBuilder {
+	b.testInstallBase = testInstallBase
+	return b
 }
 
-func AutoGenRustBenchmarkConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
+func (b *MaybeAutoGenTestConfigBuilder) SetDeviceTemplate(deviceTemplate string) *MaybeAutoGenTestConfigBuilder {
+	b.deviceTemplate = deviceTemplate
+	return b
+}
+
+func (b *MaybeAutoGenTestConfigBuilder) SetHostTemplate(hostTemplate string) *MaybeAutoGenTestConfigBuilder {
+	b.hostTemplate = hostTemplate
+	return b
+}
+
+func (b *MaybeAutoGenTestConfigBuilder) SetHostUnitTestTemplate(hostUnitTestTemplate string) *MaybeAutoGenTestConfigBuilder {
+	b.hostUnitTestTemplate = hostUnitTestTemplate
+	return b
+}
+
+func (b *MaybeAutoGenTestConfigBuilder) Build() android.Path {
+	config := append(b.config, b.configsForAutogenerated...)
+	path, autogenPath := testConfigPath(b.ctx, b.testConfigProp, b.testSuites, b.autoGenConfig, b.testConfigTemplateProp)
 	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
+		templatePath := getTestConfigTemplate(b.ctx, b.testConfigTemplateProp)
 		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), config, "")
+			autogenTemplate(b.ctx, b.name, autogenPath, templatePath.String(), config, b.outputFileName, b.testInstallBase)
 		} else {
-			if ctx.Device() {
-				autogenTemplate(ctx, autogenPath, "${RustDeviceBenchmarkConfigTemplate}", config, "")
+			if b.ctx.Device() {
+				autogenTemplate(b.ctx, b.name, autogenPath, b.deviceTemplate, config, b.outputFileName, b.testInstallBase)
 			} else {
-				autogenTemplate(ctx, autogenPath, "${RustHostBenchmarkConfigTemplate}", config, "")
+				if Bool(b.unitTest) {
+					autogenTemplate(b.ctx, b.name, autogenPath, b.hostUnitTestTemplate, config, b.outputFileName, b.testInstallBase)
+				} else {
+					autogenTemplate(b.ctx, b.name, autogenPath, b.hostTemplate, config, b.outputFileName, b.testInstallBase)
+				}
 			}
 		}
 		return autogenPath
 	}
+	if len(b.configsForAutogenerated) > 0 {
+		b.ctx.ModuleErrorf("Extra tradefed configurations were provided for an autogenerated xml file, but the autogenerated xml file was not used.")
+	}
 	return path
 }
 
-func AutoGenRobolectricTestConfig(ctx android.ModuleContext, testConfigProp *string, testConfigTemplateProp *string,
-	testSuites []string, autoGenConfig *bool) android.Path {
-	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
-	if autogenPath != nil {
-		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
-		if templatePath.Valid() {
-			autogenTemplate(ctx, autogenPath, templatePath.String(), nil, "")
-		} else {
-			autogenTemplate(ctx, autogenPath, "${RobolectricTestConfigTemplate}", nil, "")
-		}
-		return autogenPath
+func autogenTemplate(ctx android.ModuleContext, name string, output android.WritablePath, template string, configs []Config, outputFileName string, testInstallBase string) {
+	if template == "" {
+		ctx.ModuleErrorf("Empty template")
 	}
-	return path
+	var configStrings []string
+	for _, config := range configs {
+		configStrings = append(configStrings, config.Config())
+	}
+	extraConfigs := strings.Join(configStrings, fmt.Sprintf("\\n%s", test_xml_indent))
+	extraConfigs = proptools.NinjaAndShellEscape(extraConfigs)
+
+	ctx.Build(pctx, android.BuildParams{
+		Rule:        autogenTestConfig,
+		Description: "test config",
+		Output:      output,
+		Args: map[string]string{
+			"name":            name,
+			"template":        template,
+			"extraConfigs":    extraConfigs,
+			"outputFileName":  outputFileName,
+			"testInstallBase": testInstallBase,
+		},
+	})
 }
 
 var autogenInstrumentationTest = pctx.StaticRule("autogenInstrumentationTest", blueprint.RuleParams{