1 /* 2 * Copyright (C) 2019 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 com.android.tradefed.config.Option; 20 import com.android.tradefed.config.OptionClass; 21 import com.android.tradefed.device.DeviceNotAvailableException; 22 import com.android.tradefed.device.ITestDevice; 23 import com.android.tradefed.device.UserInfo; 24 import com.android.tradefed.error.HarnessRuntimeException; 25 import com.android.tradefed.invoker.TestInformation; 26 import com.android.tradefed.log.LogUtil.CLog; 27 28 import java.util.Map; 29 30 /** 31 * A {@link ITargetPreparer} that switches to the specified user kind in setUp. By default it 32 * remains in the current user, and no switching is performed. 33 * 34 * <p>Tries to restore device user state by switching back to the pre-execution current user. 35 */ 36 @OptionClass(alias = "switch-user-target-preparer") 37 public class SwitchUserTargetPreparer extends BaseTargetPreparer { 38 @Option( 39 name = "user-type", 40 description = "The type of user to switch to before the module run." 41 ) 42 private UserInfo.UserType mUserToSwitchTo = UserInfo.UserType.CURRENT; 43 44 private int mPreExecutionCurrentUser; 45 46 @Override setUp(TestInformation testInformation)47 public void setUp(TestInformation testInformation) 48 throws TargetSetupError, DeviceNotAvailableException { 49 ITestDevice device = testInformation.getDevice(); 50 setUserToSwitchTo(device); 51 52 mPreExecutionCurrentUser = device.getCurrentUser(); 53 Map<Integer, UserInfo> userInfos = device.getUserInfos(); 54 55 if (userInfos 56 .get(mPreExecutionCurrentUser) 57 .isUserType(mUserToSwitchTo, mPreExecutionCurrentUser)) { 58 CLog.i( 59 "User %d is already user type %s, no action.", 60 mPreExecutionCurrentUser, mUserToSwitchTo.toString()); 61 return; 62 } 63 64 for (UserInfo userInfo : userInfos.values()) { 65 if (userInfo.isUserType(mUserToSwitchTo, mPreExecutionCurrentUser)) { 66 CLog.i( 67 "User %d is user type %s, switching from %d", 68 userInfo.userId(), mUserToSwitchTo.toString(), mPreExecutionCurrentUser); 69 if (!device.switchUser(userInfo.userId())) { 70 throw new TargetSetupError( 71 String.format("Device failed to switch to user %d", userInfo.userId()), 72 device.getDeviceDescriptor()); 73 } 74 return; 75 } 76 } 77 78 throw new TargetSetupError( 79 String.format( 80 "Failed switch to user type %s, no user of that type exists", 81 mUserToSwitchTo), 82 device.getDeviceDescriptor()); 83 } 84 85 @Override tearDown(TestInformation testInformation, Throwable e)86 public void tearDown(TestInformation testInformation, Throwable e) 87 throws DeviceNotAvailableException { 88 // Restore the previous user as the foreground. 89 if (testInformation.getDevice().switchUser(mPreExecutionCurrentUser)) { 90 CLog.d("Successfully switched back to user id: %d", mPreExecutionCurrentUser); 91 } else { 92 CLog.w("Could not switch back to the user id: %d", mPreExecutionCurrentUser); 93 } 94 } 95 96 /** 97 * In some form factors running on headless system user mode, it is restricted to switch to the 98 * {@link UserInfo.UserType#SYSTEM SYSTEM} user. In such cases, change the {@link 99 * #mUserToSwitchTo} to the {@link UserInfo.UserType#MAIN MAIN} user. 100 */ setUserToSwitchTo(ITestDevice device)101 private void setUserToSwitchTo(ITestDevice device) throws DeviceNotAvailableException { 102 try { 103 if (UserInfo.UserType.SYSTEM.equals(mUserToSwitchTo) 104 && device.isHeadlessSystemUserMode() 105 && !device.canSwitchToHeadlessSystemUser()) { 106 mUserToSwitchTo = UserInfo.UserType.MAIN; 107 CLog.i("SwitchUserTargetPreparer is configured to switch to the MAIN user."); 108 } 109 } catch (HarnessRuntimeException e) { 110 CLog.w("Unable to get the main user switch-ability. Error: ", e); 111 } 112 } 113 } 114