1 /*
2  * Copyright (C) 2016 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.aetest.tradefed.targetprep;
18 
19 import static com.android.tradefed.util.BuildTestsZipUtils.getApkFile;
20 
21 import com.android.tradefed.build.IBuildInfo;
22 import com.android.tradefed.config.Option;
23 import com.android.tradefed.config.OptionClass;
24 import com.android.tradefed.device.DeviceNotAvailableException;
25 import com.android.tradefed.device.ITestDevice;
26 import com.android.tradefed.log.LogUtil.CLog;
27 import com.android.tradefed.targetprep.AltDirBehavior;
28 import com.android.tradefed.targetprep.BaseTargetPreparer;
29 import com.android.tradefed.targetprep.ITargetCleaner;
30 import com.android.tradefed.targetprep.ITargetPreparer;
31 import com.android.tradefed.targetprep.TargetSetupError;
32 
33 import java.io.File;
34 import java.io.IOException;
35 import java.util.ArrayList;
36 import java.util.List;
37 
38 /**
39  * A {@link ITargetPreparer} that creates a managed profile in the testing device.
40  *
41  * <p>This was forked off Android Enterprise team's target preparer, where its method of resolving
42  * the test APK was substituted with the platform's implementation. It is also modified to use the
43  * standard TradeFed way of disabling a preparer. The two classes are otherwise the same.
44  */
45 @OptionClass(alias = "ae-test-create-managed-profile")
46 public class AeTestManagedProfileCreator extends BaseTargetPreparer implements ITargetCleaner {
47 
48     @Option(
49             name = "profile-owner-component",
50             description = "Profile owner component to set; optional")
51     private String mProfileOwnerComponent = null;
52 
53     @Option(name = "profile-owner-apk", description = "Profile owner apk path; optional")
54     private String mProfileOwnerApk = null;
55 
56     @Option(
57             name = "alt-dir",
58             description =
59                     "Alternate directory to look for the apk if the apk is not in the tests "
60                             + "zip file. For each alternate dir, will look in //, //data/app, "
61                             + "//DATA/app, //DATA/app/apk_name/ and //DATA/priv-app/apk_name/. "
62                             + "Can be repeated. Look for apks in last alt-dir first.")
63     private List<File> mAltDirs = new ArrayList<>();
64 
65     @Option(
66             name = "alt-dir-behavior",
67             description =
68                     "The order of alternate directory to be used "
69                             + "when searching for apks to install")
70     private AltDirBehavior mAltDirBehavior = AltDirBehavior.FALLBACK;
71 
72     /** UserID for user in managed profile. */
73     private int mManagedProfileUserId;
74 
75     /** {@inheritDoc} */
76     @Override
setUp(ITestDevice device, IBuildInfo buildInfo)77     public void setUp(ITestDevice device, IBuildInfo buildInfo)
78             throws TargetSetupError, DeviceNotAvailableException {
79         String pmCommand =
80                 "pm create-user --profileOf 0 --managed "
81                         + "TestProfile_"
82                         + System.currentTimeMillis();
83         String pmCommandOutput = device.executeShellCommand(pmCommand);
84 
85         // Extract the id of the new user.
86         String[] pmCmdTokens = pmCommandOutput.split("\\s+");
87         if (!pmCmdTokens[0].contains("Success:")) {
88             throw new TargetSetupError("Managed profile creation failed.");
89         }
90         mManagedProfileUserId = Integer.parseInt(pmCmdTokens[pmCmdTokens.length - 1]);
91 
92         // Start managed profile user.
93         device.startUser(mManagedProfileUserId);
94 
95         CLog.i(String.format("New user created: %d", mManagedProfileUserId));
96 
97         if (mProfileOwnerComponent != null && mProfileOwnerApk != null) {
98             device.installPackageForUser(
99                     getApk(device, buildInfo, mProfileOwnerApk), true, mManagedProfileUserId);
100             String command =
101                     String.format(
102                             "dpm set-profile-owner --user %d %s",
103                             mManagedProfileUserId, mProfileOwnerComponent);
104             String commandOutput = device.executeShellCommand(command);
105             String[] cmdTokens = commandOutput.split("\\s+");
106             if (!cmdTokens[0].contains("Success:")) {
107                 throw new RuntimeException(
108                         String.format(
109                                 "Fail to setup %s as profile owner.", mProfileOwnerComponent));
110             }
111 
112             CLog.i(
113                     String.format(
114                             "%s was set as profile owner of user %d",
115                             mProfileOwnerComponent, mManagedProfileUserId));
116         }
117 
118         // Reboot device to create the apps in managed profile.
119         device.reboot();
120         device.waitForDeviceAvailable();
121     }
122 
123     /** {@inheritDoc} */
124     @Override
tearDown(ITestDevice testDevice, IBuildInfo buildInfo, Throwable throwable)125     public void tearDown(ITestDevice testDevice, IBuildInfo buildInfo, Throwable throwable)
126             throws DeviceNotAvailableException {
127         testDevice.removeUser(mManagedProfileUserId);
128     }
129 
130     /**
131      * Get the {@link File} object that refers to the APK to install.
132      *
133      * <p>This is a replacement of the method with a similar signature in the original class's super
134      * class, so that the above code can align as closely to the original class as possible.
135      */
getApk(ITestDevice device, IBuildInfo buildInfo, String fileName)136     private File getApk(ITestDevice device, IBuildInfo buildInfo, String fileName)
137             throws TargetSetupError {
138         try {
139             File apkFile =
140                     getApkFile(
141                             buildInfo,
142                             mProfileOwnerApk,
143                             mAltDirs,
144                             mAltDirBehavior,
145                             false /* use resource as fallback */,
146                             null /* device signing key */);
147             if (apkFile == null) {
148                 throw new TargetSetupError(
149                         String.format("Test app %s was not found.", mProfileOwnerApk),
150                         device.getDeviceDescriptor());
151             }
152             return apkFile;
153         } catch (IOException e) {
154             throw new TargetSetupError(
155                     String.format(
156                             "failed to resolve apk path for apk %s in build %s",
157                             mProfileOwnerApk, buildInfo.toString()),
158                     e,
159                     device.getDeviceDescriptor());
160         }
161     }
162 }
163