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 com.android.annotations.VisibleForTesting; 20 import com.android.tradefed.device.DeviceNotAvailableException; 21 import com.android.tradefed.device.ITestDevice; 22 import com.android.tradefed.device.UserInfo; 23 import com.android.tradefed.invoker.TestInformation; 24 import com.android.tradefed.log.LogUtil.CLog; 25 26 import com.google.common.base.Strings; 27 28 import java.util.ArrayList; 29 import java.util.Map; 30 31 // Not directly unit tested, but its clients are 32 public final class UserHelper { 33 34 private static final String TF_CREATED_USER = "tf_created_user"; 35 36 @VisibleForTesting static final String USER_SETUP_COMPLETE = "user_setup_complete"; 37 38 /** System property used to indicate which Android user is running the test. */ 39 public static final String RUN_TESTS_AS_USER_KEY = "RUN_TESTS_AS_USER"; 40 createUser(ITestDevice device, boolean reuseTestUser)41 public static int createUser(ITestDevice device, boolean reuseTestUser) 42 throws DeviceNotAvailableException, TargetSetupError { 43 if (reuseTestUser) { 44 Integer existingTFUser = findExistingTradefedUser(device); 45 if (existingTFUser != null) { 46 return existingTFUser; 47 } 48 } 49 50 cleanupOldUsersIfLimitReached(device); 51 52 try { 53 int userId = device.createUser(TF_CREATED_USER); 54 CLog.d("Marking user %d as setup complete", userId); 55 device.setSetting(userId, "secure", USER_SETUP_COMPLETE, "1"); 56 return userId; 57 } catch (IllegalStateException e) { 58 throw new TargetSetupError("Failed to create user.", e, device.getDeviceDescriptor()); 59 } 60 } 61 62 /** 63 * Gets the user id to run the tests as, from the {@link #RUN_TESTS_AS_USER_KEY} property. 64 * 65 * <p>If the property is not set or invalid, returns the current user. 66 */ getRunTestsAsUser(TestInformation testInfo)67 public static int getRunTestsAsUser(TestInformation testInfo) 68 throws DeviceNotAvailableException { 69 ITestDevice device = testInfo.getDevice(); 70 String val = testInfo.properties().get(RUN_TESTS_AS_USER_KEY); 71 if (!Strings.isNullOrEmpty(val)) { 72 try { 73 return Integer.parseInt(val); 74 } catch (Exception e) { 75 CLog.e("Failed to parse the userId for " + RUN_TESTS_AS_USER_KEY + " due to " + e); 76 } 77 } 78 79 // Fall back to the current user. 80 return device.getCurrentUser(); 81 } 82 cleanupOldUsersIfLimitReached(ITestDevice device)83 private static void cleanupOldUsersIfLimitReached(ITestDevice device) 84 throws DeviceNotAvailableException { 85 ArrayList<Integer> tfCreatedUsers = new ArrayList<>(); 86 int existingUsersCount = 0; 87 for (Map.Entry<Integer, UserInfo> entry : device.getUserInfos().entrySet()) { 88 UserInfo userInfo = entry.getValue(); 89 String userName = userInfo.userName(); 90 91 if (!userInfo.isGuest()) { 92 // Guest users don't fall under the quota. 93 existingUsersCount++; 94 } 95 if (userName != null && userName.equals(TF_CREATED_USER)) { 96 tfCreatedUsers.add(entry.getKey()); 97 } 98 } 99 100 if (existingUsersCount >= device.getMaxNumberOfUsersSupported()) { 101 // Reached the maximum number of users allowed. Remove stale users to free up space. 102 for (int userId : tfCreatedUsers) { 103 device.removeUser(userId); 104 } 105 } 106 } 107 findExistingTradefedUser(ITestDevice device)108 private static Integer findExistingTradefedUser(ITestDevice device) 109 throws DeviceNotAvailableException { 110 for (Map.Entry<Integer, UserInfo> entry : device.getUserInfos().entrySet()) { 111 String userName = entry.getValue().userName(); 112 113 if (userName != null && userName.equals(TF_CREATED_USER)) { 114 return entry.getKey(); 115 } 116 } 117 return null; 118 } 119 UserHelper()120 private UserHelper() { 121 throw new UnsupportedOperationException("provide only static methods"); 122 } 123 } 124