From a26516c575d99c518f8eb687a35410da576e305e Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Wed, 30 Jun 2021 08:52:54 -0700
Subject: [PATCH] Verifying success of setting default Launcher in tests

Test: local runs
Bug: 187080582
Change-Id: Iece234e58e2faa80996001f6de1943f065522400
---
 .../launcher3/util/rule/ShellCommandRule.java | 58 +++++++++++++++----
 1 file changed, 48 insertions(+), 10 deletions(-)

diff --git a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
index 0ec0f026d9..2b2fef4ae6 100644
--- a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
@@ -15,21 +15,27 @@
  */
 package com.android.launcher3.util.rule;
 
-import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
-
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
+
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
-
-import org.junit.rules.TestRule;
-import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
+import android.util.Log;
 
 import androidx.annotation.Nullable;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.UiDevice;
 
+import com.android.systemui.shared.system.PackageManagerWrapper;
+
+import org.junit.Assert;
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.util.ArrayList;
+
 /**
  * Test rule which executes a shell command at the start of the test.
  */
@@ -37,10 +43,19 @@ public class ShellCommandRule implements TestRule {
 
     private final String mCmd;
     private final String mRevertCommand;
+    private final boolean mCheckSuccess;
+    private final Runnable mAdditionalChecks;
 
-    public ShellCommandRule(String cmd, @Nullable String revertCommand) {
+    public ShellCommandRule(String cmd, @Nullable String revertCommand, boolean checkSuccess,
+            Runnable additionalChecks) {
         mCmd = cmd;
         mRevertCommand = revertCommand;
+        mCheckSuccess = checkSuccess;
+        mAdditionalChecks = additionalChecks;
+    }
+
+    public ShellCommandRule(String cmd, @Nullable String revertCommand) {
+        this(cmd, revertCommand, false, null);
     }
 
     @Override
@@ -48,12 +63,27 @@ public class ShellCommandRule implements TestRule {
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
-                UiDevice.getInstance(getInstrumentation()).executeShellCommand(mCmd);
+                final String result =
+                        UiDevice.getInstance(getInstrumentation()).executeShellCommand(mCmd);
+                if (mCheckSuccess) {
+                    Assert.assertTrue(
+                            "Failed command: " + mCmd + ", result: " + result,
+                            "Success".equals(result.replaceAll("\\s", "")));
+                }
+                if (mAdditionalChecks != null) mAdditionalChecks.run();
                 try {
                     base.evaluate();
                 } finally {
                     if (mRevertCommand != null) {
-                        UiDevice.getInstance(getInstrumentation()).executeShellCommand(mRevertCommand);
+                        final String revertResult = UiDevice.getInstance(
+                                getInstrumentation()).executeShellCommand(
+                                mRevertCommand);
+                        if (mCheckSuccess) {
+                            Assert.assertTrue(
+                                    "Failed command: " + mRevertCommand
+                                            + ", result: " + revertResult,
+                                    "Success".equals(result.replaceAll("\\s", "")));
+                        }
                     }
                 }
             }
@@ -72,7 +102,15 @@ public class ShellCommandRule implements TestRule {
      * Sets the target launcher as default launcher.
      */
     public static ShellCommandRule setDefaultLauncher() {
-        return new ShellCommandRule(getLauncherCommand(getLauncherInMyProcess()), null);
+        final ActivityInfo launcher = getLauncherInMyProcess();
+        Log.d("b/187080582", "Launcher: " + new ComponentName(launcher.packageName, launcher.name)
+                .flattenToString());
+        return new ShellCommandRule(getLauncherCommand(launcher), null, true, () ->
+                Assert.assertEquals("Setting default launcher failed",
+                        new ComponentName(launcher.packageName, launcher.name)
+                                .flattenToString(),
+                        PackageManagerWrapper.getInstance().getHomeActivities(new ArrayList<>())
+                                .flattenToString()));
     }
 
     public static String getLauncherCommand(ActivityInfo launcher) {