Make every UI test generate screenshots for local test runs

Change-Id: Idb7962e62c0c95f8a50792e9342562c6d8b6ba42

2 files changed, 37 insertions, 55 deletions

diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index b6878bba6..ca6d968cc 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -64,8 +64,12 @@ import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 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;
@@ -89,23 +93,22 @@ public abstract class AbstractLauncherUiTest {
     public static final long DEFAULT_UI_TIMEOUT = 10000;
     protected static final int LONG_WAIT_TIME_MS = 60000;
     private static final String TAG = "AbstractLauncherUiTest";
+    private static int sScreenshotCount = 0;
 
     protected MainThreadExecutor mMainThreadExecutor = new MainThreadExecutor();
-    protected final UiDevice mDevice;
+    protected final UiDevice mDevice = UiDevice.getInstance(getInstrumentation());
     protected final LauncherInstrumentation mLauncher;
     protected Context mTargetContext;
     protected String mTargetPackage;
 
     protected AbstractLauncherUiTest() {
-        final Instrumentation instrumentation = getInstrumentation();
-        mDevice = UiDevice.getInstance(instrumentation);
         try {
             mDevice.setOrientationNatural();
         } catch (RemoteException e) {
             throw new RuntimeException(e);
         }
         if (TestHelpers.isInLauncherProcess()) Utilities.enableRunningInTestHarnessForTests();
-        mLauncher = new LauncherInstrumentation(instrumentation);
+        mLauncher = new LauncherInstrumentation(getInstrumentation());
     }
 
     @Rule
@@ -163,6 +166,36 @@ public abstract class AbstractLauncherUiTest {
                 }
             } : 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);
+            }
+        }
+
+        @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();
+            mDevice.takeScreenshot(new File(pathname));
+        }
+    };
+
     @Before
     public void setUp() throws Exception {
         mTargetContext = InstrumentationRegistry.getTargetContext();
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 7578dfff2..d4cfe3aa2 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -24,11 +24,8 @@ import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-import android.util.Log;
-
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
-import androidx.test.uiautomator.UiDevice;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherState;
@@ -45,62 +42,14 @@ import com.android.launcher3.widget.WidgetsFullSheet;
 import com.android.launcher3.widget.WidgetsRecyclerView;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
-import org.junit.rules.TestWatcher;
-import org.junit.runner.Description;
 import org.junit.runner.RunWith;
 
-import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.IOException;
-
 @LargeTest
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3 extends AbstractLauncherUiTest {
-    private static final String TAG = "TaplTestsAosp";
     private static final String APP_NAME = "LauncherTestApp";
 
-    private static int sScreenshotCount = 0;
-
-    public static class FailureWatcher extends TestWatcher {
-        private UiDevice mDevice;
-
-        public FailureWatcher(UiDevice device) {
-            this.mDevice = device;
-        }
-
-        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);
-            }
-        }
-
-        @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();
-            mDevice.takeScreenshot(new File(pathname));
-        }
-    }
-
-    @Rule
-    public TestWatcher mFailureWatcher = new FailureWatcher(mDevice);
-
     @Before
     public void setUp() throws Exception {
         super.setUp();