From ba2226164c9479afd1a0dae290fb2b0535aa8e97 Mon Sep 17 00:00:00 2001
From: Ivan Lozano <ivanlozano@google.com>
Date: Tue, 21 Sep 2021 10:49:13 -0400
Subject: [PATCH] rust: Add require_root and vendor install base

Adds the require_root property to rust_test modules so that tests can
run as root. Also adds an install base to allow vendor tests to install
to a different path, otherwise installs to 'unrestricted' like cc_test.

Bug: 200602232
Test: test with require_root: true runs as root
Test: test installed to new base directory
Change-Id: Ifa716871ff87471f0c0caff7af558d47db85ccd1
---
 rust/test.go        | 23 +++++++++++++++++++++--
 tradefed/autogen.go |  8 ++++----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/rust/test.go b/rust/test.go
index e95b47cff..56da509b5 100644
--- a/rust/test.go
+++ b/rust/test.go
@@ -59,6 +59,10 @@ type TestProperties struct {
 
 	// Test options.
 	Test_options TestOptions
+
+	// Add RootTargetPreparer to auto generated test config. This guarantees the test to run
+	// with root permission.
+	Require_root *bool
 }
 
 // A test module is a binary module with extra --test compiler flag
@@ -109,12 +113,27 @@ func (test *testDecorator) compilerProps() []interface{} {
 }
 
 func (test *testDecorator) install(ctx ModuleContext) {
+	testInstallBase := "/data/local/tests/unrestricted"
+	if ctx.RustModule().InVendor() || ctx.RustModule().UseVndk() {
+		testInstallBase = "/data/local/tests/vendor"
+	}
+
+	var configs []tradefed.Config
+	if Bool(test.Properties.Require_root) {
+		configs = append(configs, tradefed.Object{"target_preparer", "com.android.tradefed.targetprep.RootTargetPreparer", nil})
+	} else {
+		var options []tradefed.Option
+		options = append(options, tradefed.Option{Name: "force-root", Value: "false"})
+		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,
-		nil,
-		test.Properties.Auto_gen_config)
+		configs,
+		test.Properties.Auto_gen_config,
+		testInstallBase)
 
 	dataSrcPaths := android.PathsForModuleSrc(ctx, test.Properties.Data)
 
diff --git a/tradefed/autogen.go b/tradefed/autogen.go
index 3d96c8457..da5582973 100644
--- a/tradefed/autogen.go
+++ b/tradefed/autogen.go
@@ -227,17 +227,17 @@ func AutoGenPythonBinaryHostTestConfig(ctx android.ModuleContext, testConfigProp
 }
 
 func AutoGenRustTestConfig(ctx android.ModuleContext, testConfigProp *string,
-	testConfigTemplateProp *string, testSuites []string, config []Config, autoGenConfig *bool) android.Path {
+	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, "")
+			autogenTemplate(ctx, autogenPath, templatePath.String(), config, testInstallBase)
 		} else {
 			if ctx.Device() {
-				autogenTemplate(ctx, autogenPath, "${RustDeviceTestConfigTemplate}", config, "")
+				autogenTemplate(ctx, autogenPath, "${RustDeviceTestConfigTemplate}", config, testInstallBase)
 			} else {
-				autogenTemplate(ctx, autogenPath, "${RustHostTestConfigTemplate}", config, "")
+				autogenTemplate(ctx, autogenPath, "${RustHostTestConfigTemplate}", config, testInstallBase)
 			}
 		}
 		return autogenPath