Fix UserManagerTest waitForUserRemoval timeout.

import android.annotation.UserIdInt;

import android.app.ActivityManager;

import android.content.Context;

import android.content.Intent;

import android.content.pm.PackageManager;

import android.content.pm.UserInfo;

import android.content.pm.UserProperties;

import androidx.test.filters.SmallTest;

import androidx.test.runner.AndroidJUnit4;

import com.android.compatibility.common.util.BlockingBroadcastReceiver;

import com.google.common.collect.ImmutableList;

import com.google.common.collect.Range;

import java.util.concurrent.Executors;

import java.util.concurrent.TimeUnit;

import java.util.concurrent.atomic.AtomicInteger;

import java.util.function.Function;

import java.util.stream.Collectors;

/** Test {@link UserManager} functionality. */

/**

 * Test {@link UserManager} functionality.

 *

 *  atest com.android.server.pm.UserManagerTest

 */

@Postsubmit

@RunWith(AndroidJUnit4.class)

public final class UserManagerTest {

    private PackageManager mPackageManager;

    private ArraySet<Integer> mUsersToRemove;

    private UserSwitchWaiter mUserSwitchWaiter;

    private UserRemovalWaiter mUserRemovalWaiter;

    private int mOriginalCurrentUserId;

    @Before

    public void setUp() throws Exception {

        mOriginalCurrentUserId = ActivityManager.getCurrentUser();

        mUserManager = UserManager.get(mContext);

        mActivityManager = mContext.getSystemService(ActivityManager.class);

        mPackageManager = mContext.getPackageManager();

        mUserSwitchWaiter = new UserSwitchWaiter(TAG, SWITCH_USER_TIMEOUT_SECONDS);

        mUserRemovalWaiter = new UserRemovalWaiter(mContext, TAG, REMOVE_USER_TIMEOUT_SECONDS);

        mUsersToRemove = new ArraySet<>();

        removeExistingUsers();

    @After

    public void tearDown() throws Exception {

        if (mOriginalCurrentUserId != ActivityManager.getCurrentUser()) {

            switchUser(mOriginalCurrentUserId);

        }

        mUserSwitchWaiter.close();

        // Making a copy of mUsersToRemove to avoid ConcurrentModificationException

        mUsersToRemove.stream().toList().forEach(this::removeUser);

        mUserRemovalWaiter.close();

    }

    private void removeExistingUsers() {

        mUserManager.setUserRestriction(UserManager.DISALLOW_REMOVE_USER, /* value= */ true,

                asHandle(currentUser));

        try {

            runThenWaitForUserRemoval(() -> {

            assertThat(mUserManager.removeUserWhenPossible(user1.getUserHandle(),

                    /* overrideDevicePolicy= */ true))

                    .isEqualTo(UserManager.REMOVE_RESULT_REMOVED);

            }, user1.id); // wait for user removal

            waitForUserRemoval(user1.id);

        } finally {

            mUserManager.setUserRestriction(UserManager.DISALLOW_REMOVE_USER, /* value= */ false,

                    asHandle(currentUser));

        assertThat(hasUser(user1.id)).isTrue();

        assertThat(getUser(user1.id).isEphemeral()).isTrue();

        runThenWaitForUserRemoval(() -> {

        // Switch back to the starting user.

        switchUser(startUser);

        // User will be removed once switch is complete

        }, user1.id); // wait for user removal

        waitForUserRemoval(user1.id);

        assertThat(hasUser(user1.id)).isFalse();

    }

        // Switch to the user just created.

        switchUser(testUser.id);

        runThenWaitForUserRemoval(() -> {

        switchUserThenRun(startUser, () -> {

            // While the switch is happening, call removeUserWhenPossible for the current user.

            assertThat(mUserManager.removeUserWhenPossible(testUser.getUserHandle(),

            assertThat(getUser(testUser.id).isEphemeral()).isTrue();

        }); // wait for user switch - startUser

        // User will be removed once switch is complete

        }, testUser.id); // wait for user removal

        waitForUserRemoval(testUser.id);

        assertThat(hasUser(testUser.id)).isFalse();

    }

            assertThat(getUser(testUser.id).isEphemeral()).isTrue();

        }); // wait for user switch - testUser

        runThenWaitForUserRemoval(() -> {

        // Switch back to the starting user.

        switchUser(startUser);

        // User will be removed once switch is complete

        }, testUser.id); // wait for user removal

        waitForUserRemoval(testUser.id);

        assertThat(hasUser(testUser.id)).isFalse();

    }

    public void testRemoveUserWhenPossible_nonCurrentUserRemoved() throws Exception {

        final UserInfo user1 = createUser("User 1", /* flags= */ 0);

        runThenWaitForUserRemoval(() -> {

        assertThat(mUserManager.removeUserWhenPossible(user1.getUserHandle(),

                /* overrideDevicePolicy= */ false))

                .isEqualTo(UserManager.REMOVE_RESULT_REMOVED);

        }, user1.id); // wait for user removal

        waitForUserRemoval(user1.id);

        assertThat(hasUser(user1.id)).isFalse();

    }

                UserManager.USER_TYPE_PROFILE_MANAGED,

                parentUser.id);

        runThenWaitForUserRemoval(() -> {

        assertThat(mUserManager.removeUserWhenPossible(parentUser.getUserHandle(),

                /* overrideDevicePolicy= */ false))

                .isEqualTo(UserManager.REMOVE_RESULT_REMOVED);

        }, parentUser.id); // wait for user removal

        waitForUserRemoval(parentUser.id);

        assertThat(hasUser(parentUser.id)).isFalse();

        assertThat(hasUser(cloneProfileUser.id)).isFalse();

    @Nullable

    private UserInfo getUser(int id) {

        List<UserInfo> list = mUserManager.getUsers();

        List<UserInfo> list = mUserManager.getAliveUsers();

        for (UserInfo user : list) {

            if (user.id == id) {

    @MediumTest

    @Test

    public void testConcurrentUserCreate() throws Exception {

        int userCount = mUserManager.getUserCount();

        int userCount = mUserManager.getAliveUsers().size();

        int maxSupportedUsers = UserManager.getMaxSupportedUsers();

        int canBeCreatedCount = maxSupportedUsers - userCount;

        // Test exceeding the limit while running in parallel

        }

        es.shutdown();

        es.awaitTermination(20, TimeUnit.SECONDS);

        assertThat(mUserManager.getUserCount()).isEqualTo(maxSupportedUsers);

        assertThat(mUserManager.getAliveUsers().size()).isEqualTo(maxSupportedUsers);

        assertThat(created.get()).isEqualTo(canBeCreatedCount);

    }

    }

    private void removeUser(UserHandle userHandle) {

        runThenWaitForUserRemoval(

                () -> mUserManager.removeUser(userHandle),

                userHandle == null ? UserHandle.USER_NULL : userHandle.getIdentifier()

        );

        mUserManager.removeUser(userHandle);

        waitForUserRemoval(userHandle.getIdentifier());

    }

    private void removeUser(int userId) {

        runThenWaitForUserRemoval(

                () -> mUserManager.removeUser(userId),

                userId

        );

    }

    private void runThenWaitForUserRemoval(Runnable runnable, int userIdToWaitUntilDeleted) {

        if (!hasUser(userIdToWaitUntilDeleted)) {

            runnable.run();

            mUsersToRemove.remove(userIdToWaitUntilDeleted);

            return;

        mUserManager.removeUser(userId);

        waitForUserRemoval(userId);

    }

        Function<Intent, Boolean> checker = intent -> {

            UserHandle userHandle = intent.getParcelableExtra(Intent.EXTRA_USER, UserHandle.class);

            return userHandle != null && userHandle.getIdentifier() == userIdToWaitUntilDeleted;

        };

        BlockingBroadcastReceiver blockingBroadcastReceiver = BlockingBroadcastReceiver.create(

                mContext, Intent.ACTION_USER_REMOVED, checker);

        blockingBroadcastReceiver.setTimeout(REMOVE_USER_TIMEOUT_SECONDS);

        blockingBroadcastReceiver.register();

        try (blockingBroadcastReceiver) {

            runnable.run();

        }

        mUsersToRemove.remove(userIdToWaitUntilDeleted);

    private void waitForUserRemoval(int userId) {

        mUserRemovalWaiter.waitFor(userId);

        mUsersToRemove.remove(userId);

    }

    private UserInfo createUser(String name, int flags) {

/*

 * Copyright (C) 2023 The Android Open Source Project

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.android.server.pm;

import static org.junit.Assert.fail;

import android.content.BroadcastReceiver;

import android.content.Context;

import android.content.Intent;

import android.content.IntentFilter;

import android.os.UserHandle;

import android.os.UserManager;

import android.util.Log;

import java.io.Closeable;

import java.io.IOException;

import java.util.Map;

import java.util.concurrent.ConcurrentHashMap;

import java.util.concurrent.CountDownLatch;

import java.util.concurrent.TimeUnit;

public class UserRemovalWaiter extends BroadcastReceiver implements Closeable {

    private final Context mContext;

    private final UserManager mUserManager;

    private final String mTag;

    private final long mTimeoutMillis;

    private final Map<Integer, CountDownLatch> mMap = new ConcurrentHashMap<>();

    private CountDownLatch getLatch(final int userId) {

        return mMap.computeIfAbsent(userId, absentKey -> new CountDownLatch(1));

    }

    public UserRemovalWaiter(Context context, String tag, int timeoutInSeconds) {

        mContext = context;

        mUserManager = UserManager.get(mContext);

        mTag = tag;

        mTimeoutMillis = timeoutInSeconds * 1000L;

        mContext.registerReceiver(this, new IntentFilter(Intent.ACTION_USER_REMOVED));

    }

    @Override

    public void close() throws IOException {

        mContext.unregisterReceiver(this);

    }

    @Override

    public void onReceive(Context context, Intent intent) {

        if (Intent.ACTION_USER_REMOVED.equals(intent.getAction())) {

            int userId = intent.getIntExtra(Intent.EXTRA_USER_HANDLE, UserHandle.USER_NULL);

            Log.i(mTag, "ACTION_USER_REMOVED received for user " + userId);

            getLatch(userId).countDown();

        }

    }

    /**

     * Waits for the removal of the given user, or fails if it times out.

     */

    public void waitFor(int userId) {

        Log.i(mTag, "Waiting for user " + userId + " to be removed");

        CountDownLatch latch = getLatch(userId);

        long startTime = System.currentTimeMillis();

        while (System.currentTimeMillis() - startTime < mTimeoutMillis) {

            if (hasUserGone(userId) || waitLatchForOneSecond(latch)) {

                Log.i(mTag, "User " + userId + " is removed in "

                        + (System.currentTimeMillis() - startTime) + " ms");

                return;

            }

        }

        fail("Timeout waiting for user removal. userId = " + userId);

    }

    private boolean hasUserGone(int userId) {

        return mUserManager.getAliveUsers().stream().noneMatch(x -> x.id == userId);

    }

    private boolean waitLatchForOneSecond(CountDownLatch latch) {

        try {

            return latch.await(1, TimeUnit.SECONDS);

        } catch (InterruptedException e) {

            Log.e(mTag, "Thread interrupted unexpectedly.", e);

            return false;

        }

    }

}