1 /*
2  * Copyright (C) 2023 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.tradefed.targetprep;
18 
19 import static com.android.tradefed.log.LogUtil.CLog;
20 
21 import static java.util.stream.Collectors.joining;
22 
23 import com.android.tradefed.config.Option;
24 import com.android.tradefed.config.OptionClass;
25 import com.android.tradefed.invoker.TestInformation;
26 import com.android.tradefed.log.ITestLogger;
27 import com.android.tradefed.result.ITestLoggerReceiver;
28 import com.android.tradefed.result.error.DeviceErrorIdentifier;
29 import com.android.tradefed.result.error.InfraErrorIdentifier;
30 import com.android.tradefed.util.CommandResult;
31 import com.android.tradefed.util.DeviceActionUtil;
32 import com.android.tradefed.util.DeviceActionUtil.Command;
33 
34 import com.google.common.annotations.VisibleForTesting;
35 import com.google.common.collect.ImmutableList;
36 
37 import java.io.File;
38 import java.io.IOException;
39 import java.util.ArrayList;
40 import java.util.Collection;
41 import java.util.HashSet;
42 import java.util.Set;
43 
44 import javax.annotation.Nullable;
45 
46 /** A {@link ITargetPreparer} to perform device actions. */
47 @OptionClass(alias = "device-action")
48 public class DeviceActionTargetPreparer extends BaseTargetPreparer implements ITestLoggerReceiver {
49     private static final String REPEATED_FLAG_DELIMITER = ", ";
50     private static final String FILE_ARG_FORMAT = "file_%s=%s";
51     private static final String MAINLINE_MODULES_TAG = "mainline_modules";
52     private static final String APKS_ZIPS_TAG = "apks_zips";
53     private static final String TRAIN_FOLDER_TAG = "train_folder";
54     private static final String ENABLE_ROLLBACK_FLAG = "enable_rollback";
55     private static final String DEV_KEY_SIGNED_FLAG = "dev_key_signed";
56 
57     // Common options for all commands.
58     @Option(name = "da-command", description = "Command name of device action.", mandatory = true)
59     private Command mCommand;
60 
61     @Option(
62             name = "bundletool-jar",
63             description = "The file name of the bundletool jar.",
64             mandatory = true)
65     File mBundletoolJar;
66 
67     @Option(
68             name = "device-action-jar",
69             description = "The file name of the device action jar.",
70             mandatory = true)
71     File mDeviceActionJar;
72 
73     // Options for install mainline command.
74     @Option(name = "enable-rollback", description = "Enable rollback if the mainline update fails.")
75     private boolean mEnableRollback = true;
76 
77     @Option(name = "dev-key-signed", description = "If the modules are signed by dev keys.")
78     private boolean mDevKeySigned = false;
79 
80     @Option(name = "train-folder", description = "The path of the train folder.")
81     private File mTrainFolderPath;
82 
83     @Option(
84             name = "mainline-modules",
85             description =
86                     "The name of module file to be installed on device. Can be repeated. The file"
87                             + " can be apk/apex/apks. If one file is apks, then all files should be"
88                             + " apks.")
89     private Set<File> mModuleFiles = new HashSet<>();
90 
91     @Option(
92             name = "apks-zips",
93             description = "The name of zip file containing train apks. Can be repeated.")
94     private Set<File> mZipFiles = new HashSet<>();
95 
96     private DeviceActionUtil mDeviceActionUtil;
97 
98     private ITestLogger mTestLogger;
99 
100     @Override
setUp(TestInformation testInfo)101     public void setUp(TestInformation testInfo) throws TargetSetupError {
102         if (mDeviceActionUtil == null) {
103             try {
104                 mDeviceActionUtil = DeviceActionUtil.create(mDeviceActionJar, mBundletoolJar);
105             } catch (DeviceActionUtil.DeviceActionConfigError e) {
106                 throw new TargetSetupError("Failed to create DeviceActionUtil", e, e.getErrorId());
107             }
108         }
109 
110         CommandResult result = performAction(testInfo);
111         if (result == null) {
112             CLog.w("No action performed.");
113             return;
114         }
115         CLog.i(result.getStdout());
116         CLog.i(result.getStderr());
117 
118         TargetSetupError error = null;
119         try {
120             mDeviceActionUtil.generateLogFile(result);
121             if (mTestLogger != null) {
122                 mDeviceActionUtil.saveToLogs(mCommand, mTestLogger);
123             }
124         } catch (IOException e) {
125             error =
126                     new TargetSetupError(
127                             "Failed to save log files.",
128                             e,
129                             InfraErrorIdentifier.FAIL_TO_CREATE_FILE);
130         }
131         if (result.getExitCode() != 0) {
132             error =
133                     new TargetSetupError(
134                             String.format(
135                                     "Failed to execute device action command %s. The status is %s.",
136                                     mCommand, result.getStatus()),
137                             DeviceErrorIdentifier.DEVICE_ACTION_EXECUTION_FAILURE);
138         } else {
139             CLog.i("Device action completed.");
140         }
141 
142         if (error != null) {
143             throw error;
144         }
145     }
146 
147     @Override
setTestLogger(ITestLogger testLogger)148     public void setTestLogger(ITestLogger testLogger) {
149         mTestLogger = testLogger;
150     }
151 
152     @VisibleForTesting
setDeviceActionUtil(DeviceActionUtil deviceActionUtil)153     void setDeviceActionUtil(DeviceActionUtil deviceActionUtil) {
154         mDeviceActionUtil = deviceActionUtil;
155     }
156 
157     /** Performs device action. Returns null if no action performed. */
158     @Nullable
performAction(TestInformation testInfo)159     private CommandResult performAction(TestInformation testInfo) {
160         String deviceId = testInfo.getDevice().getSerialNumber();
161         switch (mCommand) {
162             case RESET:
163                 return reset(deviceId);
164             case INSTALL_MAINLINE:
165                 return installMainline(deviceId);
166             default:
167                 return null;
168         }
169     }
170 
reset(String deviceId)171     private CommandResult reset(String deviceId) {
172         return mDeviceActionUtil.execute(Command.RESET, deviceId, ImmutableList.of());
173     }
174 
175     @Nullable
installMainline(String deviceId)176     private CommandResult installMainline(String deviceId) {
177         ArrayList<String> args = new ArrayList<>();
178         if (!mZipFiles.isEmpty()) {
179             args.add(createFileArg(APKS_ZIPS_TAG, mZipFiles));
180         } else if (!mModuleFiles.isEmpty()) {
181             args.add(createFileArg(MAINLINE_MODULES_TAG, mModuleFiles));
182         } else if (mTrainFolderPath != null && mTrainFolderPath.exists()) {
183             args.add(createFileArg(TRAIN_FOLDER_TAG, mTrainFolderPath));
184         } else {
185             CLog.i("No module provided. No mainline update.");
186             return null;
187         }
188         if (mEnableRollback) {
189             args.add(ENABLE_ROLLBACK_FLAG);
190         }
191         if (mDevKeySigned) {
192             args.add(DEV_KEY_SIGNED_FLAG);
193         }
194         return mDeviceActionUtil.execute(Command.INSTALL_MAINLINE, deviceId, args);
195     }
196 
createFileArg(String tag, Collection<File> files)197     private static String createFileArg(String tag, Collection<File> files) {
198         return files.stream()
199                 .map(file -> createFileArg(tag, file))
200                 .collect(joining(REPEATED_FLAG_DELIMITER));
201     }
202 
createFileArg(String tag, File file)203     private static String createFileArg(String tag, File file) {
204         return String.format(FILE_ARG_FORMAT, tag, file.getAbsolutePath());
205     }
206 }
207