From ad15dd719ee50d57a9000cd6ceabfb8a63dd3773 Mon Sep 17 00:00:00 2001
From: Sebastian Franco <fransebas@google.com>
Date: Tue, 23 Apr 2024 10:35:55 -0700
Subject: [PATCH] Make sure we keep the screen record even if the test succeeds

Bug: 322823478
Test: testEdu
Test: testDevicePressMenu
Flag: NA
Change-Id: Ia11d3cbc0a10e15d716b1845e3590bca2bd1ff76
---
 .../launcher3/ui/TaplTestsLauncher3Test.java      |  4 +++-
 .../android/launcher3/ui/TaplWorkProfileTest.java |  2 ++
 .../launcher3/util/rule/ScreenRecordRule.java     | 15 ++++++++++++++-
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
index 342eedf2be..4a67600024 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
@@ -21,8 +21,9 @@ import static org.junit.Assert.assertNotNull;
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.Launcher;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
+import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,6 +32,7 @@ import org.junit.runner.RunWith;
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3Test extends AbstractLauncherUiTest<Launcher> {
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testDevicePressMenu() throws Exception {
diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
index e08d37cb69..70a3dd001a 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
@@ -45,6 +45,7 @@ import com.android.launcher3.allapps.WorkPausedCard;
 import com.android.launcher3.allapps.WorkProfileManager;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
@@ -196,6 +197,7 @@ public class TaplWorkProfileTest extends AbstractLauncherUiTest<Launcher> {
 
     }
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testEdu() {
diff --git a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
index 7a5cf2c2cd..ff96afb51d 100644
--- a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
@@ -49,6 +49,9 @@ public class ScreenRecordRule implements TestRule {
             return base;
         }
 
+        final Boolean keepRecordOnSuccess = description.getAnnotation(KeepRecordOnSuccess.class)
+                != null;
+
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
@@ -70,7 +73,7 @@ public class ScreenRecordRule implements TestRule {
                     device.executeShellCommand("kill -INT " + screenRecordPid);
                     Log.e(TAG, "Screenrecord captured at: " + outputFile);
                     output.close();
-                    if (success) {
+                    if (success && !keepRecordOnSuccess) {
                         automation.executeShellCommand("rm " + outputFile);
                     }
                 }
@@ -85,4 +88,14 @@ public class ScreenRecordRule implements TestRule {
     @Target(ElementType.METHOD)
     public @interface ScreenRecord {
     }
+
+
+    /**
+     * Interface to indicate that we should keep the screen record even if the test succeeds, use
+     * sparingly since it uses a lof of memory.
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface KeepRecordOnSuccess {
+    }
 }