UserChecker: add option to cleanup users

Bug: 135057192
Bug: 130047592
Test: unit tests
Test: Ran cts test with --user-type secondary and --user-cleanup. User
  was switch-back and deleted at the end.
Test: did above AND manually created user mid-test. User was deleted at
  end.
Change-Id: I76f6f408ff19c67dea09bf4a49abce339edde161
Merged-In: I76f6f408ff19c67dea09bf4a49abce339edde161

2 files changed, 96 insertions, 2 deletions

diff --git a/src/com/android/tradefed/suite/checker/UserChecker.java b/src/com/android/tradefed/suite/checker/UserChecker.java
index 92b12dd2e..e90019155 100644
--- a/src/com/android/tradefed/suite/checker/UserChecker.java
+++ b/src/com/android/tradefed/suite/checker/UserChecker.java
@@ -41,6 +41,18 @@ public class UserChecker implements ISystemStatusChecker {
     )
     private UserInfo.UserType mUserToSwitchTo = UserInfo.UserType.CURRENT;
 
+    @Option(
+        name = "user-cleanup",
+        description =
+                "If true, attempt to cleanup any changes made to users:"
+                        + "\n - switch to previous current-user"
+                        + "\n - remove any created users"
+                        + "\n\nThis does NOT:"
+                        + "\n - attempt to re-create a user that was deleted"
+                        + "\n - start/stop existing users if their running status changed"
+    )
+    private boolean mCleanup = false;
+
     private UserInfo mPreCurrentUserInfo = null;
     private Map<Integer, UserInfo> mPreUsersInfo = null;
     private int mSwitchedToUserId = -1;
@@ -98,7 +110,15 @@ public class UserChecker implements ISystemStatusChecker {
                                 "User %d was the currentUser before, has changed to %d",
                                 mPreCurrentUserInfo.userId(), postCurrentUserInfo.userId()));
             }
-            // TODO(b/130047592): do cleanup
+            if (mCleanup) {
+                if (!device.switchUser(mPreCurrentUserInfo.userId())) {
+                    errors.add(
+                            String.format(
+                                    "Failed to switch back to previous current user %d."
+                                            + " Check if it was removed.",
+                                    mPreCurrentUserInfo.userId()));
+                }
+            }
         }
 
         for (UserInfo preUserInfo : mPreUsersInfo.values()) {
@@ -123,7 +143,11 @@ public class UserChecker implements ISystemStatusChecker {
                             String.format(
                                     "User %d was created during test and not deleted", postUserId));
                 }
-                // TODO(b/130047592): do cleanup
+                if (mCleanup) {
+                    if (!device.removeUser(postUserId)) {
+                        errors.add(String.format("Failed to remove new user %d", postUserId));
+                    }
+                }
             }
         }
 
diff --git a/tests/src/com/android/tradefed/suite/checker/UserCheckerTest.java b/tests/src/com/android/tradefed/suite/checker/UserCheckerTest.java
index d4fece268..c6f0ca594 100644
--- a/tests/src/com/android/tradefed/suite/checker/UserCheckerTest.java
+++ b/tests/src/com/android/tradefed/suite/checker/UserCheckerTest.java
@@ -20,6 +20,7 @@ import java.util.Map;
 import java.util.HashMap;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.ITestDevice;
@@ -114,6 +115,75 @@ public class UserCheckerTest {
     }
 
     @Test
+    public void testCreateCleanup() throws Exception {
+        UserChecker checker = new UserChecker();
+        OptionSetter mOptionSetter = new OptionSetter(checker);
+        mOptionSetter.setOptionValue("user-type", "secondary");
+        mOptionSetter.setOptionValue("user-cleanup", "true");
+        ITestDevice preDevice =
+                mockDeviceUserState(
+                        /* currentUser=  */ 0,
+                        /* userIds=        */ new Integer[] {0},
+                        /* flags=        */ new Integer[] {0},
+                        /* isRunning= */ new Boolean[] {true});
+        when(preDevice.createUser("Tfsecondary", false, false)).thenReturn(10);
+        when(preDevice.switchUser(10)).thenReturn(true);
+
+        assertEquals(CheckStatus.SUCCESS, checker.preExecutionCheck(preDevice).getStatus());
+        verify(preDevice, times(1)).createUser("Tfsecondary", false, false);
+        verify(preDevice, times(1)).switchUser(10);
+
+        ITestDevice postDevice =
+                mockDeviceUserState(
+                        /* currentUser=  */ 10,
+                        /* userIds=        */ new Integer[] {0, 10},
+                        /* flags=        */ new Integer[] {0, 0},
+                        /* isRunning= */ new Boolean[] {true, true});
+        when(postDevice.switchUser(0)).thenReturn(true);
+        when(postDevice.removeUser(10)).thenReturn(true);
+        assertEquals(CheckStatus.SUCCESS, checker.postExecutionCheck(postDevice).getStatus());
+        verify(postDevice, times(1)).switchUser(0);
+        verify(postDevice, times(1)).removeUser(10);
+    }
+
+    @Test
+    public void testCreateCleanup_cleanupFail() throws Exception {
+        UserChecker checker = new UserChecker();
+        OptionSetter mOptionSetter = new OptionSetter(checker);
+        mOptionSetter.setOptionValue("user-type", "secondary");
+        mOptionSetter.setOptionValue("user-cleanup", "true");
+        ITestDevice preDevice =
+                mockDeviceUserState(
+                        /* currentUser=  */ 0,
+                        /* userIds=        */ new Integer[] {0},
+                        /* flags=        */ new Integer[] {0},
+                        /* isRunning= */ new Boolean[] {true});
+        when(preDevice.createUser("Tfsecondary", false, false)).thenReturn(10);
+        when(preDevice.switchUser(10)).thenReturn(true);
+
+        assertEquals(CheckStatus.SUCCESS, checker.preExecutionCheck(preDevice).getStatus());
+        verify(preDevice, times(1)).createUser("Tfsecondary", false, false);
+        verify(preDevice, times(1)).switchUser(10);
+
+        ITestDevice postDevice =
+                mockDeviceUserState(
+                        /* currentUser=  */ 10,
+                        /* userIds=        */ new Integer[] {0, 10},
+                        /* flags=        */ new Integer[] {0, 0},
+                        /* isRunning= */ new Boolean[] {true, true});
+        when(postDevice.switchUser(0)).thenReturn(false);
+        when(postDevice.removeUser(10)).thenReturn(false);
+        StatusCheckerResult result = checker.postExecutionCheck(postDevice);
+        verify(postDevice, times(1)).switchUser(0);
+        verify(postDevice, times(1)).removeUser(10);
+        assertEquals(CheckStatus.FAILED, result.getStatus());
+        assertTrue(
+                result.getErrorMessage()
+                        .contains("Failed to switch back to previous current user 0"));
+        assertTrue(result.getErrorMessage().contains("Failed to remove new user 10"));
+    }
+
+    @Test
     /** Returns FAILED in the precessense of errors */
     public void testAllErrorsIsFailed() throws Exception {
         UserChecker checker = new UserChecker();