diff --git a/rust/benchmark.go b/rust/benchmark.go
index 0e84243..b417a2d 100644
--- a/rust/benchmark.go
+++ b/rust/benchmark.go
@@ -112,12 +112,14 @@
 }
 
 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 0cc3bca..ecc7d5d 100644
--- a/rust/test.go
+++ b/rust/test.go
@@ -130,13 +130,16 @@
 		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)