From cf17e849bbd3da6c75f532feb3bc497aeac471cd Mon Sep 17 00:00:00 2001
From: Yuexi Ma <yuexima@google.com>
Date: Wed, 28 Apr 2021 22:46:53 -0700
Subject: Add an option to take screenshot after app launch

Test: atest csuite-harness-tests
Bug: 184659711
Change-Id: I8a80dddee70208f794eea861d7a94e2985e6798b
---
 .../compatibility/testtype/AppLaunchTest.java      | 20 ++++++++++++++++
 .../compatibility/testtype/AppLaunchTestTest.java  | 27 +++++++++++++++++++++-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/harness/src/main/java/com/android/compatibility/testtype/AppLaunchTest.java b/harness/src/main/java/com/android/compatibility/testtype/AppLaunchTest.java
index e15f960..7a52735 100644
--- a/harness/src/main/java/com/android/compatibility/testtype/AppLaunchTest.java
+++ b/harness/src/main/java/com/android/compatibility/testtype/AppLaunchTest.java
@@ -59,6 +59,12 @@ import java.util.Set;
 /** A test that verifies that a single app can be successfully launched. */
 public class AppLaunchTest
         implements IDeviceTest, IRemoteTest, IConfigurationReceiver, ITestFilterReceiver {
+    @VisibleForTesting static final String SCREENSHOT_AFTER_LAUNCH = "screenshot-after-launch";
+
+    @Option(
+            name = SCREENSHOT_AFTER_LAUNCH,
+            description = "Whether to take a screenshost after a package is launched.")
+    private boolean mScreenshotAfterLaunch;
 
     @Option(name = "package-name", description = "Package name of testing app.")
     private String mPackageName;
@@ -204,6 +210,20 @@ public class AppLaunchTest
                     break;
                 }
             }
+
+            if (mScreenshotAfterLaunch) {
+                try (InputStreamSource screenSource = mDevice.getScreenshot()) {
+                    listener.testLog(
+                            mPackageName + "_screenshot_" + mDevice.getSerialNumber(),
+                            LogDataType.PNG,
+                            screenSource);
+                } catch (DeviceNotAvailableException e) {
+                    CLog.e(
+                            "Device %s became unavailable while capturing screenshot, %s",
+                            mDevice.getSerialNumber(), e.toString());
+                    throw e;
+                }
+            }
         } finally {
             reportResult(listener, testDescription, result);
             stopPackage();
diff --git a/harness/src/test/java/com/android/compatibility/testtype/AppLaunchTestTest.java b/harness/src/test/java/com/android/compatibility/testtype/AppLaunchTestTest.java
index 070cf85..0c23cc3 100644
--- a/harness/src/test/java/com/android/compatibility/testtype/AppLaunchTestTest.java
+++ b/harness/src/test/java/com/android/compatibility/testtype/AppLaunchTestTest.java
@@ -15,16 +15,20 @@
  */
 package com.android.compatibility.testtype;
 
+import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.metrics.proto.MetricMeasurement.Metric;
+import com.android.tradefed.result.FileInputStreamSource;
 import com.android.tradefed.result.ITestInvocationListener;
+import com.android.tradefed.result.InputStreamSource;
 import com.android.tradefed.result.TestDescription;
 import com.android.tradefed.testtype.InstrumentationTest;
 import com.android.tradefed.util.CommandResult;
 import com.android.tradefed.util.CommandStatus;
 
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -38,12 +42,14 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertThrows;
 
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.InOrder;
+import org.mockito.Mockito;
 
 import java.util.HashMap;
 import java.util.HashSet;
@@ -56,6 +62,7 @@ public final class AppLaunchTestTest {
     private final ITestInvocationListener mMockListener = mock(ITestInvocationListener.class);
     private static final String TEST_PACKAGE_NAME = "package_name";
     private static final TestInformation NULL_TEST_INFORMATION = null;
+    @Rule public final TemporaryFolder tempFolder = new TemporaryFolder();
 
     @Test
     public void run_testFailed() throws DeviceNotAvailableException {
@@ -77,6 +84,24 @@ public final class AppLaunchTestTest {
         verifyPassedAndEndedCall(mMockListener);
     }
 
+    @Test
+    public void run_takeScreenShot_savesToTestLog() throws Exception {
+        InstrumentationTest instrumentationTest = createPassingInstrumentationTest();
+        AppLaunchTest appLaunchTest = createLaunchTestWithInstrumentation(instrumentationTest);
+        new OptionSetter(appLaunchTest)
+                .setOptionValue(AppLaunchTest.SCREENSHOT_AFTER_LAUNCH, "true");
+        ITestDevice mMockDevice = mock(ITestDevice.class);
+        appLaunchTest.setDevice(mMockDevice);
+        InputStreamSource screenshotData = new FileInputStreamSource(tempFolder.newFile());
+        when(mMockDevice.getScreenshot()).thenReturn(screenshotData);
+        when(mMockDevice.getSerialNumber()).thenReturn("SERIAL");
+
+        appLaunchTest.run(NULL_TEST_INFORMATION, mMockListener);
+
+        Mockito.verify(mMockListener, times(1))
+                .testLog(Mockito.contains("screenshot"), Mockito.any(), Mockito.eq(screenshotData));
+    }
+
     @Test
     public void run_packageResetSuccess() throws DeviceNotAvailableException {
         ITestDevice mMockDevice = mock(ITestDevice.class);
-- 
cgit v1.2.3