From 3a4b79bf7b838a74918484d4fbcab58420d6da3c Mon Sep 17 00:00:00 2001
From: Himanshu Gupta <himanshuz@google.com>
Date: Fri, 7 Jun 2024 18:29:42 +0100
Subject: [PATCH] Scrolling down to locate PS after lock/unlock event.

This change fixes TAPL tests by:
1. Scrolling down to locate PS, in case the scrollbar moves after lock/unlock.
2. Retrying lock/unlock, as sometimes, the request is cancelled by UserManager (if the profile is already in that state)

Bug: 345556016
Test: atest TaplPrivateSpaceTest
Flag: NONE Tapl fix tweak.
Change-Id: Ic0cc3259a2f92065a699d694c47f65c5f68934b8
---
 .../quickstep/TaplPrivateSpaceTest.java       | 30 ++++++++++++++++---
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
index 82905085fe..23a29f7840 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
@@ -35,7 +35,6 @@ import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ScreenRecordRule;
 
 import org.junit.After;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -50,6 +49,7 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
 
     private static final String PRIVATE_PROFILE_NAME = "LauncherPrivateProfile";
     private static final String INSTALLED_APP_NAME = "Aardwolf";
+    private static final int MAX_STATE_TOGGLE_TRIES = 2;
     private static final String TAG = "TaplPrivateSpaceTest";
 
     @Override
@@ -167,14 +167,15 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
 
     @Test
     @ScreenRecordRule.ScreenRecord // b/334946529
-    @Ignore("b/339179262")
     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);
+        togglePrivateSpaceWithRetry(PrivateProfileManager.STATE_DISABLED, homeAllApps);
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
 
         homeAllApps.freeze();
         try {
@@ -186,7 +187,9 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
         }
 
         // Enable Private Space
-        togglePrivateSpace(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+        togglePrivateSpaceWithRetry(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
 
         homeAllApps.freeze();
         try {
@@ -219,6 +222,25 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
         waitForLauncherUIUpdate();
     }
 
+    private void togglePrivateSpaceWithRetry(int state, HomeAllApps homeAllApps) {
+        int togglePsCount = 0;
+        boolean shouldRetry;
+        do {
+            togglePsCount ++;
+            try {
+                togglePrivateSpace(state, homeAllApps);
+                // No need to retry if the toggle was successful.
+                shouldRetry = false;
+            } catch (AssertionError error) {
+                if (togglePsCount < MAX_STATE_TOGGLE_TRIES) {
+                    shouldRetry = true;
+                } else {
+                    throw error;
+                }
+            }
+        } while (shouldRetry);
+    }
+
     private void waitForPrivateSpaceSetup() {
         waitForLauncherCondition("Private Profile not setup",
                 launcher -> launcher.getAppsView().hasPrivateProfile(),