Merge "Keep ViewCaptureRule logic self-contained." into udc-dev
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 97e34c5..dbe4402 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -116,12 +116,11 @@
             Utilities.enableRunningInTestHarnessForTests();
         }
 
-        final ViewCaptureRule viewCaptureRule = new ViewCaptureRule();
         mOrderSensitiveRules = RuleChain
                 .outerRule(new SamplerRule())
                 .around(new NavigationModeSwitchRule(mLauncher))
-                .around(viewCaptureRule)
-                .around(new FailureWatcher(mDevice, mLauncher, viewCaptureRule.getViewCapture()));
+                .around(new ViewCaptureRule())
+                .around(new FailureWatcher(mDevice, mLauncher));
 
         mOtherLauncherActivity = context.getPackageManager().queryIntentActivities(
                 getHomeIntentInPackage(context),
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index d7c4ae3..5bd28d8 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -216,11 +216,10 @@
     }
 
     protected TestRule getRulesInsideActivityMonitor() {
-        final ViewCaptureRule viewCaptureRule = new ViewCaptureRule();
         final RuleChain inner = RuleChain
                 .outerRule(new PortraitLandscapeRunner(this))
-                .around(viewCaptureRule)
-                .around(new FailureWatcher(mDevice, mLauncher, viewCaptureRule.getViewCapture()));
+                .around(new ViewCaptureRule())
+                .around(new FailureWatcher(mDevice, mLauncher));
 
         return TestHelpers.isInLauncherProcess()
                 ? RuleChain.outerRule(ShellCommandRule.setDefaultLauncher()).around(inner)
diff --git a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
index 7ca6a06..6b11fd6 100644
--- a/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
+++ b/tests/src/com/android/launcher3/util/rule/FailureWatcher.java
@@ -6,12 +6,8 @@
 import android.os.ParcelFileDescriptor.AutoCloseInputStream;
 import android.util.Log;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.test.core.app.ApplicationProvider;
 import androidx.test.uiautomator.UiDevice;
 
-import com.android.app.viewcapture.ViewCapture;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 
@@ -32,14 +28,10 @@
     private static boolean sSavedBugreport = false;
     final private UiDevice mDevice;
     private final LauncherInstrumentation mLauncher;
-    @NonNull
-    private final ViewCapture mViewCapture;
 
-    public FailureWatcher(UiDevice device, LauncherInstrumentation launcher,
-            @NonNull ViewCapture viewCapture) {
+    public FailureWatcher(UiDevice device, LauncherInstrumentation launcher) {
         mDevice = device;
         mLauncher = launcher;
-        mViewCapture = viewCapture;
     }
 
     @Override
@@ -71,7 +63,7 @@
 
     @Override
     protected void failed(Throwable e, Description description) {
-        onError(mLauncher, description, e, mViewCapture);
+        onError(mLauncher, description, e);
     }
 
     static File diagFile(Description description, String prefix, String ext) {
@@ -82,12 +74,6 @@
 
     public static void onError(LauncherInstrumentation launcher, Description description,
             Throwable e) {
-        onError(launcher, description, e, null);
-    }
-
-    private static void onError(LauncherInstrumentation launcher, Description description,
-            Throwable e, @Nullable ViewCapture viewCapture) {
-
         final File sceenshot = diagFile(description, "TestScreenshot", "png");
         final File hierarchy = diagFile(description, "Hierarchy", "zip");
 
@@ -102,12 +88,6 @@
             out.putNextEntry(new ZipEntry("visible_windows.zip"));
             dumpCommand("cmd window dump-visible-window-views", out);
             out.closeEntry();
-
-            if (viewCapture != null) {
-                out.putNextEntry(new ZipEntry("FS/data/misc/wmtrace/failed_test.vc"));
-                viewCapture.dumpTo(out, ApplicationProvider.getApplicationContext());
-                out.closeEntry();
-            }
         } catch (Exception ignored) {
         }
 
diff --git a/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt b/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
index 0c65539..f3fff35 100644
--- a/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
+++ b/tests/src/com/android/launcher3/util/rule/ViewCaptureRule.kt
@@ -19,62 +19,101 @@
 import android.app.Application
 import android.media.permission.SafeCloseable
 import android.os.Bundle
+import android.util.Log
+import androidx.annotation.AnyThread
 import androidx.test.core.app.ApplicationProvider
 import com.android.app.viewcapture.SimpleViewCapture
 import com.android.app.viewcapture.ViewCapture.MAIN_EXECUTOR
 import com.android.launcher3.util.ActivityLifecycleCallbacksAdapter
-import org.junit.rules.TestRule
+import java.io.File
+import java.io.FileOutputStream
+import java.util.zip.ZipEntry
+import java.util.zip.ZipOutputStream
+import org.junit.rules.TestWatcher
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
 
+private const val TAG = "ViewCaptureRule"
+
 /**
  * This JUnit TestRule registers a listener for activity lifecycle events to attach a ViewCapture
  * instance that other test rules use to dump the timelapse hierarchy upon an error during a test.
  *
  * This rule will not work in OOP tests that don't have access to the activity under test.
  */
-class ViewCaptureRule : TestRule {
-    val viewCapture = SimpleViewCapture("test-view-capture")
+class ViewCaptureRule : TestWatcher() {
+    private val viewCapture = SimpleViewCapture("test-view-capture")
+    private val windowListenerCloseables = mutableListOf<SafeCloseable>()
 
     override fun apply(base: Statement, description: Description): Statement {
+        val testWatcherStatement = super.apply(base, description)
+
         return object : Statement() {
             override fun evaluate() {
-                val windowListenerCloseables = mutableListOf<SafeCloseable>()
-
-                val lifecycleCallbacks =
-                    object : ActivityLifecycleCallbacksAdapter {
-                        override fun onActivityCreated(activity: Activity, bundle: Bundle?) {
-                            super.onActivityCreated(activity, bundle)
-                            windowListenerCloseables.add(
-                                viewCapture.startCapture(
-                                    activity.window.decorView,
-                                    "${description.testClass?.simpleName}.${description.methodName}"
-                                )
-                            )
-                        }
-
-                        override fun onActivityDestroyed(activity: Activity) {
-                            super.onActivityDestroyed(activity)
-                            viewCapture.stopCapture(activity.window.decorView)
-                        }
+                val lifecycleCallbacks = createLifecycleCallbacks(description)
+                with(ApplicationProvider.getApplicationContext<Application>()) {
+                    registerActivityLifecycleCallbacks(lifecycleCallbacks)
+                    try {
+                        testWatcherStatement.evaluate()
+                    } finally {
+                        unregisterActivityLifecycleCallbacks(lifecycleCallbacks)
                     }
-
-                val application = ApplicationProvider.getApplicationContext<Application>()
-                application.registerActivityLifecycleCallbacks(lifecycleCallbacks)
-
-                try {
-                    base.evaluate()
-                } finally {
-                    application.unregisterActivityLifecycleCallbacks(lifecycleCallbacks)
-
-                    // Clean up ViewCapture references here rather than in onActivityDestroyed so
-                    // test code can access view hierarchy capture. onActivityDestroyed would delete
-                    // view capture data before FailureWatcher could output it as a test artifact.
-                    // This is on the main thread to avoid a race condition where the onDrawListener
-                    // is removed while onDraw is running, resulting in an IllegalStateException.
-                    MAIN_EXECUTOR.execute { windowListenerCloseables.onEach(SafeCloseable::close) }
                 }
             }
         }
     }
+
+    private fun createLifecycleCallbacks(description: Description) =
+        object : ActivityLifecycleCallbacksAdapter {
+            override fun onActivityCreated(activity: Activity, bundle: Bundle?) {
+                super.onActivityCreated(activity, bundle)
+                windowListenerCloseables.add(
+                    viewCapture.startCapture(
+                        activity.window.decorView,
+                        "${description.testClass?.simpleName}.${description.methodName}"
+                    )
+                )
+            }
+
+            override fun onActivityDestroyed(activity: Activity) {
+                super.onActivityDestroyed(activity)
+                viewCapture.stopCapture(activity.window.decorView)
+            }
+        }
+
+    override fun succeeded(description: Description) = cleanup()
+
+    /** If the test fails, this function will output the ViewCapture information. */
+    override fun failed(e: Throwable, description: Description) {
+        super.failed(e, description)
+
+        val testName = "${description.testClass.simpleName}.${description.methodName}"
+        val application: Application = ApplicationProvider.getApplicationContext()
+        val zip = File(application.filesDir, "ViewCapture-$testName.zip")
+
+        ZipOutputStream(FileOutputStream(zip)).use {
+            it.putNextEntry(ZipEntry("FS/data/misc/wmtrace/failed_test.vc"))
+            viewCapture.dumpTo(it, ApplicationProvider.getApplicationContext())
+            it.closeEntry()
+        }
+        cleanup()
+
+        Log.d(
+            TAG,
+            "Failed $testName due to ${e::class.java.simpleName}.\n" +
+                "\tUse go/web-hv to open dump file: \n\t\t${zip.absolutePath}"
+        )
+    }
+
+    /**
+     * Clean up ViewCapture references can't happen in onActivityDestroyed otherwise view
+     * hierarchies would be erased before they could be outputted.
+     *
+     * This is on the main thread to avoid a race condition where the onDrawListener is removed
+     * while onDraw is running, resulting in an IllegalStateException.
+     */
+    @AnyThread
+    private fun cleanup() {
+        MAIN_EXECUTOR.execute { windowListenerCloseables.onEach(SafeCloseable::close) }
+    }
 }