Use java host unit tests template for unit tests

Test: make aoa-helper-tests
Bug: 172961860
Change-Id: I2eaee277961f29adfd22e7c65248c9403bd69b81
diff --git a/java/java.go b/java/java.go
index a23c649..76f9ba2 100644
--- a/java/java.go
+++ b/java/java.go
@@ -2325,7 +2325,7 @@
 
 func (j *Test) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	j.testConfig = tradefed.AutoGenJavaTestConfig(ctx, j.testProperties.Test_config, j.testProperties.Test_config_template,
-		j.testProperties.Test_suites, j.testProperties.Auto_gen_config)
+		j.testProperties.Test_suites, j.testProperties.Auto_gen_config, j.testProperties.Test_options.Unit_test)
 
 	j.data = android.PathsForModuleSrc(ctx, j.testProperties.Data)
 
@@ -2344,7 +2344,7 @@
 
 func (j *JavaTestImport) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	j.testConfig = tradefed.AutoGenJavaTestConfig(ctx, j.prebuiltTestProperties.Test_config, nil,
-		j.prebuiltTestProperties.Test_suites, nil)
+		j.prebuiltTestProperties.Test_suites, nil, nil)
 
 	j.Import.GenerateAndroidBuildActions(ctx)
 }
diff --git a/tradefed/autogen.go b/tradefed/autogen.go
index b35f831..27d71e8 100644
--- a/tradefed/autogen.go
+++ b/tradefed/autogen.go
@@ -188,7 +188,7 @@
 }
 
 func AutoGenJavaTestConfig(ctx android.ModuleContext, testConfigProp *string, testConfigTemplateProp *string,
-	testSuites []string, autoGenConfig *bool) android.Path {
+	testSuites []string, autoGenConfig *bool, unitTest *bool) android.Path {
 	path, autogenPath := testConfigPath(ctx, testConfigProp, testSuites, autoGenConfig, testConfigTemplateProp)
 	if autogenPath != nil {
 		templatePath := getTestConfigTemplate(ctx, testConfigTemplateProp)
@@ -198,7 +198,11 @@
 			if ctx.Device() {
 				autogenTemplate(ctx, autogenPath, "${JavaTestConfigTemplate}", nil, "")
 			} else {
-				autogenTemplate(ctx, autogenPath, "${JavaHostTestConfigTemplate}", nil, "")
+				if Bool(unitTest) {
+					autogenTemplate(ctx, autogenPath, "${JavaHostUnitTestConfigTemplate}", nil, "")
+				} else {
+					autogenTemplate(ctx, autogenPath, "${JavaHostTestConfigTemplate}", nil, "")
+				}
 			}
 		}
 		return autogenPath
diff --git a/tradefed/config.go b/tradefed/config.go
index f7e8349..f3566a8 100644
--- a/tradefed/config.go
+++ b/tradefed/config.go
@@ -27,6 +27,7 @@
 	pctx.SourcePathVariable("InstrumentationTestConfigTemplate", "build/make/core/instrumentation_test_config_template.xml")
 	pctx.SourcePathVariable("JavaTestConfigTemplate", "build/make/core/java_test_config_template.xml")
 	pctx.SourcePathVariable("JavaHostTestConfigTemplate", "build/make/core/java_host_test_config_template.xml")
+	pctx.SourcePathVariable("JavaHostUnitTestConfigTemplate", "build/make/core/java_host_unit_test_config_template.xml")
 	pctx.SourcePathVariable("NativeBenchmarkTestConfigTemplate", "build/make/core/native_benchmark_test_config_template.xml")
 	pctx.SourcePathVariable("NativeHostTestConfigTemplate", "build/make/core/native_host_test_config_template.xml")
 	pctx.SourcePathVariable("NativeTestConfigTemplate", "build/make/core/native_test_config_template.xml")