Snap for 9883365 from c356139bb561403bf8574553a7dbc35e8aee4452 to tm-qpr3-release

Change-Id: Ia465f237e486ff9cb03a9ad0f7825ffe310d023b

3 files changed, 46 insertions, 5 deletions

diff --git a/libraries/health/rules/src/android/platform/test/rule/ArtifactSaver.java b/libraries/health/rules/src/android/platform/test/rule/ArtifactSaver.java
index b8ee0e5e1..eec75273f 100644
--- a/libraries/health/rules/src/android/platform/test/rule/ArtifactSaver.java
+++ b/libraries/health/rules/src/android/platform/test/rule/ArtifactSaver.java
@@ -16,8 +16,11 @@
 
 package android.platform.test.rule;
 
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
 import android.os.ParcelFileDescriptor;
 
+import androidx.annotation.NonNull;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.UiDevice;
 
@@ -50,10 +53,15 @@ public class ArtifactSaver {
             // Can happen when the description is from a ClassRule
             suffix = "EntireClassExecution";
         }
+        Class<?> testClass = description.getTestClass();
+
+        // Can have null class if this is a synthetic suite
+        String className = testClass != null ? testClass.getSimpleName() : "SUITE";
         return artifactFile(
-                prefix
+                "TestScreenshot-"
+                        + prefix
                         + "-"
-                        + description.getTestClass().getSimpleName()
+                        + className
                         + "."
                         + suffix
                         + "."
@@ -61,11 +69,10 @@ public class ArtifactSaver {
     }
 
     public static void onError(Description description, Throwable e) {
-        final UiDevice device = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
-        final File screenshot = artifactFile(description, "TestScreenshot", "png");
+        final UiDevice device = getUiDevice();
         final File hierarchy = artifactFile(description, "Hierarchy", "zip");
 
-        device.takeScreenshot(screenshot);
+        final File screenshot = takeDebugScreenshot(description, device, "OnFailure");
 
         // Dump accessibility hierarchy
         try {
@@ -106,6 +113,28 @@ public class ArtifactSaver {
         }
     }
 
+    private static UiDevice getUiDevice() {
+        return UiDevice.getInstance(getInstrumentation());
+    }
+
+    @NonNull
+    private static File takeDebugScreenshot(Description description, UiDevice device,
+            String prefix) {
+        final File screenshot = artifactFile(description, prefix, "png");
+        device.takeScreenshot(screenshot);
+        return screenshot;
+    }
+
+    public static void takeDebugScreenshot(Description description, String prefix) {
+        File screenshotFile = takeDebugScreenshot(description, getUiDevice(), prefix);
+        android.util.Log.e(
+                TAG,
+                "Screenshot taken in test: "
+                        + description.getMethodName()
+                        + ",\nscreenshot will be saved to "
+                        + screenshotFile);
+    }
+
     private static void dumpCommandAndOutput(String cmd, OutputStream out) throws IOException {
         out.write(("\n\n" + cmd + "\n").getBytes());
         dumpCommandOutput(cmd, out);
diff --git a/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java b/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
index 83ca09662..af1327125 100644
--- a/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
+++ b/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
@@ -34,6 +34,8 @@ import java.util.Optional;
 /** The device-side implementation of StsLogic. */
 public class StsExtraBusinessLogicTestCase extends ExtraBusinessLogicTestCase implements StsLogic {
 
+    private static final String LOG_TAG = StsExtraBusinessLogicTestCase.class.getSimpleName();
+
     private LocalDate deviceSpl = null;
     @Rule public DescriptionProvider descriptionProvider = new DescriptionProvider();
 
@@ -46,6 +48,10 @@ public class StsExtraBusinessLogicTestCase extends ExtraBusinessLogicTestCase im
         // set in test/sts/tools/sts-tradefed/res/config/sts-base-dynamic-*.xml
         String stsDynamicPlan =
                 InstrumentationRegistry.getArguments().getString("sts-dynamic-plan");
+        if (stsDynamicPlan == null) {
+            Log.w(LOG_TAG, "cannot get STS dynamic plan; this likely isn't run in STS");
+            return List.of();
+        }
         return StsLogic.getExtraBusinessLogicForPlan(stsDynamicPlan);
     }
 
diff --git a/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java b/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
index 80c5cc95b..e93166dd4 100644
--- a/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
+++ b/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
@@ -35,6 +35,8 @@ import java.util.Optional;
 public class StsExtraBusinessLogicHostTestBase extends ExtraBusinessLogicHostTestBase
         implements StsLogic {
 
+    private static final String LOG_TAG = StsExtraBusinessLogicHostTestBase.class.getSimpleName();
+
     private LocalDate deviceSpl = null;
     @Rule public DescriptionProvider descriptionProvider = new DescriptionProvider();
 
@@ -47,6 +49,10 @@ public class StsExtraBusinessLogicHostTestBase extends ExtraBusinessLogicHostTes
     public List<String> getExtraBusinessLogics() {
         // set in test/sts/tools/sts-tradefed/res/config/sts-base-dynamic-*.xml
         String stsDynamicPlan = getBuild().getBuildAttributes().get("sts-dynamic-plan");
+        if (stsDynamicPlan == null) {
+            Log.w(LOG_TAG, "cannot get STS dynamic plan; this likely isn't run in STS");
+            return List.of();
+        }
         return StsLogic.getExtraBusinessLogicForPlan(stsDynamicPlan);
     }