Merge "Adding Tapl test for locking/unlocking Private Space" into main
diff --git a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
index efe773b..c64ac23 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
@@ -27,6 +27,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 
+import com.android.launcher3.allapps.PrivateProfileManager;
 import com.android.launcher3.tapl.HomeAllApps;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.PrivateSpaceContainer;
@@ -164,6 +165,59 @@
         }
     }
 
+    @Test
+    @ScreenRecordRule.ScreenRecord // b/334946529
+    public void testPrivateSpaceLockingBehaviour() throws IOException {
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
+        HomeAllApps homeAllApps = mLauncher.getAllApps();
+
+        // Disable Private Space
+        togglePrivateSpace(PrivateProfileManager.STATE_DISABLED, homeAllApps);
+
+        homeAllApps.freeze();
+        try {
+            // Verify Locked View elements are present.
+            homeAllApps.getPrivateSpaceLockedView();
+        } finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+
+        // Enable Private Space
+        togglePrivateSpace(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+
+        homeAllApps.freeze();
+        try {
+            // Verify UnLocked View elements are present.
+            homeAllApps.getPrivateSpaceUnlockedView();
+        } finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+    }
+
+    private void togglePrivateSpace(int state, HomeAllApps homeAllApps) {
+        homeAllApps.freeze();
+        try {
+            // Try Toggling Private Space
+            homeAllApps.togglePrivateSpace();
+        }  finally {
+            // UnFreeze
+            homeAllApps.unfreeze();
+        }
+        PrivateProfileManager manager = getFromLauncher(l -> l.getAppsView()
+                .getPrivateProfileManager());
+        waitForLauncherCondition("Private profile toggle to state: " + state + " failed",
+                launcher -> {
+                    manager.reset();
+                    return manager.getCurrentState() == state;
+                },
+                LauncherInstrumentation.WAIT_TIME_MS);
+        // Wait for Launcher UI to be updated with Private Space Items.
+        waitForLauncherUIUpdate();
+    }
+
     private void waitForPrivateSpaceSetup() {
         waitForLauncherCondition("Private Profile not setup",
                 launcher -> launcher.getAppsView().hasPrivateProfile(),
diff --git a/tests/tapl/com/android/launcher3/tapl/AllApps.java b/tests/tapl/com/android/launcher3/tapl/AllApps.java
index f3ec852..71cd19d 100644
--- a/tests/tapl/com/android/launcher3/tapl/AllApps.java
+++ b/tests/tapl/com/android/launcher3/tapl/AllApps.java
@@ -58,6 +58,7 @@
     private static final String FAST_SCROLLER_RES_ID = "fast_scroller";
     private static final Pattern EVENT_ALT_ESC_UP = Pattern.compile(
             "Key event: KeyEvent.*?action=ACTION_UP.*?keyCode=KEYCODE_ESCAPE.*?metaState=0");
+    private static final String UNLOCK_BUTTON_VIEW_RES_ID = "ps_lock_unlock_button";
 
     private final int mHeight;
     private final int mIconHeight;
@@ -433,7 +434,28 @@
     public PrivateSpaceContainer getPrivateSpaceUnlockedView() {
         final UiObject2 allAppsContainer = verifyActiveContainer();
         final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
-        return new PrivateSpaceContainer(mLauncher, appListRecycler, this);
+        return new PrivateSpaceContainer(mLauncher, appListRecycler, this, true);
+    }
+
+    /** Returns PrivateSpaceContainer in locked state, if present in view. */
+    @NonNull
+    public PrivateSpaceContainer getPrivateSpaceLockedView() {
+        final UiObject2 allAppsContainer = verifyActiveContainer();
+        final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
+        return new PrivateSpaceContainer(mLauncher, appListRecycler, this, false);
+    }
+
+    /**
+     * Toggles Lock/Unlock of Private Space, changing the All Apps Ui.
+     */
+    public void togglePrivateSpace() {
+        final UiObject2 allAppsContainer = verifyActiveContainer();
+        final UiObject2 appListRecycler = getAppListRecycler(allAppsContainer);
+        UiObject2 unLockButtonView = mLauncher.waitForObjectInContainer(appListRecycler,
+                UNLOCK_BUTTON_VIEW_RES_ID);
+        mLauncher.waitForObjectEnabled(unLockButtonView, "Private Space Unlock Button");
+        mLauncher.assertTrue("PS Unlock Button is non-clickable", unLockButtonView.isClickable());
+        unLockButtonView.click();
     }
 
     protected abstract void verifyVisibleContainerOnDismiss();
diff --git a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
index 0e65ffb..c21bae3 100644
--- a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
+++ b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceContainer.java
@@ -31,16 +31,24 @@
     private final LauncherInstrumentation mLauncher;
     private final UiObject2 mAppListRecycler;
     private final AllApps mAppList;
+    private final boolean mPrivateSpaceEnabled;
 
     PrivateSpaceContainer(LauncherInstrumentation launcherInstrumentation,
-            UiObject2 appListRecycler, AllApps appList) {
+            UiObject2 appListRecycler, AllApps appList, boolean privateSpaceEnabled) {
         mLauncher = launcherInstrumentation;
         mAppListRecycler = appListRecycler;
         mAppList = appList;
+        mPrivateSpaceEnabled = privateSpaceEnabled;
 
-        verifyHeaderIsPresent();
-        verifyInstallAppButtonIsPresent();
-        verifyDividerIsPresent();
+        if (mPrivateSpaceEnabled) {
+            verifyHeaderIsPresent();
+            verifyInstallAppButtonIsPresent();
+            verifyDividerIsPresent();
+        } else {
+            verifyHeaderIsPresent();
+            verifyInstallAppButtonIsNotPresent();
+            verifyDividerIsNotPresent();
+        }
     }
 
     // Assert PS Header is in view.
@@ -48,7 +56,7 @@
     private void verifyHeaderIsPresent() {
         final UiObject2 psHeader = mLauncher.waitForObjectInContainer(mAppListRecycler,
                 PS_HEADER_RES_ID);
-        new PrivateSpaceHeader(mLauncher, psHeader, true);
+        new PrivateSpaceHeader(mLauncher, psHeader, mPrivateSpaceEnabled);
     }
 
 
@@ -57,11 +65,21 @@
         mAppList.getAppIcon(INSTALL_APP_TITLE);
     }
 
+    // Assert Install App Item is not present in view.
+    private void verifyInstallAppButtonIsNotPresent() {
+        mLauncher.waitUntilLauncherObjectGone(DIVIDER_RES_ID);
+    }
+
     // Assert Sys App Divider is present in view.
     private void verifyDividerIsPresent() {
         mLauncher.waitForObjectInContainer(mAppListRecycler, DIVIDER_RES_ID);
     }
 
+    // Assert Sys App Divider is not present in view.
+    private void verifyDividerIsNotPresent() {
+        mLauncher.waitUntilLauncherObjectGone(DIVIDER_RES_ID);
+    }
+
     /**
      * Verifies that a user installed app is present above the divider.
      */
@@ -73,4 +91,4 @@
         mLauncher.assertTrue("Installed App: " + appName + " is not above the divider",
                 iconCenter.y < dividerCenter.y);
     }
-}
+}
\ No newline at end of file
diff --git a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
index 3a7f038..cc64aae 100644
--- a/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
+++ b/tests/tapl/com/android/launcher3/tapl/PrivateSpaceHeader.java
@@ -45,7 +45,7 @@
         if (mPrivateSpaceEnabled) {
             verifyUnlockedState();
         } else {
-            mLauncher.fail("Private Space found in non enabled state");
+            verifyLockedState();
         }
     }
 
@@ -74,4 +74,23 @@
         mLauncher.assertEquals("PS lock text is incorrect", "Lock", lockText.getText());
 
     }
+
+    /** Verify Locked State elements in Private Space Header */
+    private void verifyLockedState() {
+        UiObject2 headerText = mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                PS_HEADER_TEXT_RES_ID);
+        mLauncher.assertEquals("PS Header Text is incorrect ",
+                "Private", headerText.getText());
+
+        UiObject2 unLockButtonView = mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                UNLOCK_BUTTON_VIEW_RES_ID);
+        mLauncher.waitForObjectEnabled(unLockButtonView, "Private Space Unlock Button");
+        mLauncher.assertTrue("PS Unlock Button is non-clickable", unLockButtonView.isClickable());
+
+        mLauncher.waitForObjectInContainer(mPrivateSpaceHeader,
+                LOCK_ICON_RES_ID);
+
+        mLauncher.waitUntilLauncherObjectGone(SETTINGS_BUTTON_RES_ID);
+        mLauncher.waitUntilLauncherObjectGone(LOCK_TEXT_RES_ID);
+    }
 }