diff --git a/rust/test.go b/rust/test.go
index 0cc3bca..4f922b4 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.AutoGenTestConfig(ctx, tradefed.AutoGenTestConfigOptions{
+		TestConfigProp:         test.Properties.Test_config,
+		TestConfigTemplateProp: test.Properties.Test_config_template,
+		TestSuites:             test.Properties.Test_suites,
+		Config:                 configs,
+		AutoGenConfig:          test.Properties.Auto_gen_config,
+		TestInstallBase:        testInstallBase,
+		DeviceTemplate:         "${RustDeviceTestConfigTemplate}",
+		HostTemplate:           "${RustHostTestConfigTemplate}",
+	})
 
 	dataSrcPaths := android.PathsForModuleSrc(ctx, test.Properties.Data)