diff --git a/java/robolectric.go b/java/robolectric.go
index 6e8d591..68f27b8 100644
--- a/java/robolectric.go
+++ b/java/robolectric.go
@@ -131,9 +131,14 @@
 	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.AutoGenTestConfig(ctx, tradefed.AutoGenTestConfigOptions{
+		TestConfigProp:         r.testProperties.Test_config,
+		TestConfigTemplateProp: r.testProperties.Test_config_template,
+		TestSuites:             r.testProperties.Test_suites,
+		AutoGenConfig:          r.testProperties.Auto_gen_config,
+		DeviceTemplate:         "${RobolectricTestConfigTemplate}",
+		HostTemplate:           "${RobolectricTestConfigTemplate}",
+	})
 	r.data = android.PathsForModuleSrc(ctx, r.testProperties.Data)
 
 	roboTestConfig := android.PathForModuleGen(ctx, "robolectric").