From e4c4e2043fa2bd11144b7dfeeb42fd2d8ff167e9 Mon Sep 17 00:00:00 2001
From: Brandon Dayauon <brdayauon@google.com>
Date: Wed, 31 Aug 2022 15:12:08 -0700
Subject: [PATCH] Added extra check to CLI command that was not checked for

check if work profile is indeed set up

 - after workProfile setup, added log of profile ID for each test..
 - added CLI check for "start-user"
 - Investigating where flakiness occurs because of failed
 work profile account set up or test just fails

bug: 243688989
test: Tested where there is no shell command in setup and all tests were skipped
Change-Id: I131b62c1a929c2e550be111a3840fd93c7435860
---
 .../android/launcher3/ui/WorkProfileTest.java | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index 5761abdea1..5757f1341c 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -48,6 +48,7 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
 
     private int mProfileUserId;
     private boolean mWorkProfileSetupSuccessful;
+    private final String TAG = "WorkProfileTest";
 
     @Before
     @Override
@@ -56,18 +57,17 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
         String output =
                 mDevice.executeShellCommand(
                         "pm create-user --profileOf 0 --managed TestProfile");
-        Log.d("b/203817455", "pm create-user; output: " + output);
-
-        if (output.startsWith("Success")){
-            assertTrue("Failed to create work profile", output.startsWith("Success"));
-            mWorkProfileSetupSuccessful = true;
-        } else {
-            return; // no need to setup launcher since all tests will skip.
-        }
+        // b/203817455
+        updateWorkProfileSetupSuccessful("pm create-user", output);
 
         String[] tokens = output.split("\\s+");
         mProfileUserId = Integer.parseInt(tokens[tokens.length - 1]);
-        mDevice.executeShellCommand("am start-user " + mProfileUserId);
+        output = mDevice.executeShellCommand("am start-user " + mProfileUserId);
+        updateWorkProfileSetupSuccessful("am start-user", output);
+
+        if (!mWorkProfileSetupSuccessful) {
+            return; // no need to setup launcher since all tests will skip.
+        }
 
         mDevice.pressHome();
         waitForLauncherCondition("Launcher didn't start", Objects::nonNull);
@@ -107,6 +107,7 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
     @Test
     public void workTabExists() {
         assumeTrue(mWorkProfileSetupSuccessful);
+        waitForWorkTabSetup();
         waitForLauncherCondition("Personal tab is missing",
                 launcher -> launcher.getAppsView().isPersonalTabVisible(),
                 LauncherInstrumentation.WAIT_TIME_MS);
@@ -183,4 +184,14 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
             }
         }, LauncherInstrumentation.WAIT_TIME_MS);
     }
+
+    private void updateWorkProfileSetupSuccessful(String cli, String output) {
+        Log.d(TAG, "updateWorkProfileSetupSuccessful, cli=" + cli + " " + "output=" + output);
+        if (output.startsWith("Success")) {
+            assertTrue(output, output.startsWith("Success"));
+            mWorkProfileSetupSuccessful = true;
+        } else {
+            mWorkProfileSetupSuccessful = false;
+        }
+    }
 }