Taking local screenshot before navigation state is restored

Change-Id: Id8003b17c68af5b4882bdbf5e14a142875791cbc
diff --git a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
index c77726e..cc79c9d 100644
--- a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
+++ b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
@@ -18,14 +18,14 @@
 
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 
-import org.junit.Rule;
-import org.junit.rules.TestRule;
+import org.junit.rules.RuleChain;
 
 /**
  * Base class for all instrumentation tests that deal with Quickstep.
  */
 public abstract class AbstractQuickStepTest extends AbstractLauncherUiTest {
-    @Rule
-    public TestRule mQuickstepOnOffExecutor =
-            new NavigationModeSwitchRule(mLauncher);
+    protected AbstractQuickStepTest() {
+        mOrderSensitiveRules = RuleChain.outerRule(new NavigationModeSwitchRule(mLauncher)).
+                around(mOrderSensitiveRules);
+    }
 }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index c7c36b0..17a5335 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -58,19 +58,17 @@
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.util.Wait;
+import com.android.launcher3.util.rule.FailureWatcher;
 import com.android.launcher3.util.rule.LauncherActivityRule;
 import com.android.launcher3.util.rule.ShellCommandRule;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
+import org.junit.rules.RuleChain;
 import org.junit.rules.TestRule;
-import org.junit.rules.TestWatcher;
-import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
-import java.io.ByteArrayOutputStream;
-import java.io.File;
 import java.io.IOException;
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
@@ -93,7 +91,6 @@
     public static final long SHORT_UI_TIMEOUT = 300;
     public static final long DEFAULT_UI_TIMEOUT = 10000;
     private static final String TAG = "AbstractLauncherUiTest";
-    private static int sScreenshotCount = 0;
 
     protected MainThreadExecutor mMainThreadExecutor = new MainThreadExecutor();
     protected final UiDevice mDevice = UiDevice.getInstance(getInstrumentation());
@@ -167,42 +164,11 @@
             } : base;
 
     @Rule
-    public TestWatcher mFailureWatcher = new TestWatcher() {
-        private void dumpViewHierarchy() {
-            final ByteArrayOutputStream stream = new ByteArrayOutputStream();
-            try {
-                mDevice.dumpWindowHierarchy(stream);
-                stream.flush();
-                stream.close();
-                for (String line : stream.toString().split("\\r?\\n")) {
-                    Log.e(TAG, line.trim());
-                }
-            } catch (IOException e) {
-                Log.e(TAG, "error dumping XML to logcat", e);
-            }
-        }
+    public RuleChain mOrderSensitiveRules = RuleChain.outerRule(new FailureWatcher(this));
 
-        @Override
-        protected void failed(Throwable e, Description description) {
-            if (mDevice == null) return;
-            final String pathname = getInstrumentation().getTargetContext().
-                    getFilesDir().getPath() + "/TaplTestScreenshot" + sScreenshotCount++ + ".png";
-            Log.e(TAG, "Failed test " + description.getMethodName() +
-                    ", screenshot will be saved to " + pathname +
-                    ", track trace is below, UI object dump is further below:\n" +
-                    Log.getStackTraceString(e));
-            dumpViewHierarchy();
-
-            try {
-                final String dumpsysResult = mDevice.executeShellCommand(
-                                "dumpsys activity service TouchInteractionService");
-                Log.d(TAG, "TouchInteractionService: " + dumpsysResult);
-            } catch (IOException ex) {
-            }
-
-            mDevice.takeScreenshot(new File(pathname));
-        }
-    };
+    public UiDevice getDevice() {
+        return mDevice;
+    }
 
     @Before
     public void setUp() throws Exception {
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
new file mode 100644
index 0000000..09cc98d
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -0,0 +1,59 @@
+package com.android.launcher3.util.rule;
+
+import static androidx.test.InstrumentationRegistry.getInstrumentation;
+
+import android.util.Log;
+
+import com.android.launcher3.ui.AbstractLauncherUiTest;
+
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+
+public class FailureWatcher extends TestWatcher {
+    private static final String TAG = "FailureWatcher";
+    private static int sScreenshotCount = 0;
+    private AbstractLauncherUiTest mAbstractLauncherUiTest;
+
+    public FailureWatcher(AbstractLauncherUiTest abstractLauncherUiTest) {
+        mAbstractLauncherUiTest = abstractLauncherUiTest;
+    }
+
+    private void dumpViewHierarchy() {
+        final ByteArrayOutputStream stream = new ByteArrayOutputStream();
+        try {
+            mAbstractLauncherUiTest.getDevice().dumpWindowHierarchy(stream);
+            stream.flush();
+            stream.close();
+            for (String line : stream.toString().split("\\r?\\n")) {
+                Log.e(TAG, line.trim());
+            }
+        } catch (IOException e) {
+            Log.e(TAG, "error dumping XML to logcat", e);
+        }
+    }
+
+    @Override
+    protected void failed(Throwable e, Description description) {
+        if (mAbstractLauncherUiTest.getDevice() == null) return;
+        final String pathname = getInstrumentation().getTargetContext().
+                getFilesDir().getPath() + "/TaplTestScreenshot" + sScreenshotCount++ + ".png";
+        Log.e(TAG, "Failed test " + description.getMethodName() +
+                ", screenshot will be saved to " + pathname +
+                ", track trace is below, UI object dump is further below:\n" +
+                Log.getStackTraceString(e));
+        dumpViewHierarchy();
+
+        try {
+            final String dumpsysResult = mAbstractLauncherUiTest.getDevice().executeShellCommand(
+                    "dumpsys activity service TouchInteractionService");
+            Log.d(TAG, "TouchInteractionService: " + dumpsysResult);
+        } catch (IOException ex) {
+        }
+
+        mAbstractLauncherUiTest.getDevice().takeScreenshot(new File(pathname));
+    }
+}