1 /*
2  * Copyright (C) 2017 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 package android.device.collectors;
17 
18 import android.device.collectors.annotations.OptionClass;
19 import android.graphics.Bitmap;
20 import android.os.Bundle;
21 import android.util.Log;
22 import androidx.annotation.VisibleForTesting;
23 import androidx.test.uiautomator.UiDevice;
24 
25 import org.junit.runner.Description;
26 import org.junit.runner.notification.Failure;
27 
28 import java.io.File;
29 import java.io.IOException;
30 import java.io.OutputStream;
31 import java.util.Map;
32 import java.util.HashMap;
33 
34 
35 /**
36  * A {@link BaseMetricListener} that captures screenshots when a test fails.
37  *
38  * <p>Dumping the UI XML requires external storage permission. See {@link BaseMetricListener} how to
39  * grant external storage permission, especially at install time. Only collecting a screenshot does
40  * not require storage permissions.
41  *
42  * <p>Options: -e screenshot-quality [0-100]: set screenshot image quality. Default is 75. -e
43  * include-ui-xml [true, false]: include the UI XML on failure too, if true.
44  */
45 @OptionClass(alias = "screenshot-failure-collector")
46 public class ScreenshotOnFailureCollector extends BaseMetricListener {
47 
48     public static final String DEFAULT_DIR = "run_listeners/screenshots";
49     public static final String KEY_INCLUDE_XML = "include-ui-xml";
50     public static final String KEY_QUALITY = "screenshot-quality";
51     public static final int DEFAULT_QUALITY = 75;
52     private boolean mIncludeUiXml = false;
53     private int mQuality = DEFAULT_QUALITY;
54 
55     private File mDestDir;
56     private UiDevice mDevice;
57 
58     // Tracks the test iterations to ensure that each failure gets unique filenames.
59     // Key: test description; value: number of iterations.
60     private Map<String, Integer> mTestIterations = new HashMap<String, Integer>();
61 
ScreenshotOnFailureCollector()62     public ScreenshotOnFailureCollector() {
63         super();
64     }
65 
66     /**
67      * Constructor to simulate receiving the instrumentation arguments. Should not be used except
68      * for testing.
69      */
70     @VisibleForTesting
ScreenshotOnFailureCollector(Bundle args)71     ScreenshotOnFailureCollector(Bundle args) {
72         super(args);
73     }
74 
75     @Override
onTestRunStart(DataRecord runData, Description description)76     public void onTestRunStart(DataRecord runData, Description description) {
77         Bundle args = getArgsBundle();
78         if (args.containsKey(KEY_QUALITY)) {
79             try {
80                 int quality = Integer.parseInt(args.getString(KEY_QUALITY));
81                 if (quality >= 0 && quality <= 100) {
82                     mQuality = quality;
83                 } else {
84                     Log.e(getTag(), String.format("Invalid screenshot quality: %d.", quality));
85                 }
86             } catch (Exception e) {
87                 Log.e(getTag(), "Failed to parse screenshot quality", e);
88             }
89         }
90 
91         if (args.containsKey(KEY_INCLUDE_XML)) {
92             mIncludeUiXml = Boolean.parseBoolean(args.getString(KEY_INCLUDE_XML));
93         }
94 
95         String dir = DEFAULT_DIR;
96         mDestDir = createAndEmptyDirectory(dir);
97     }
98 
99     @Override
onTestStart(DataRecord testData, Description description)100     public void onTestStart(DataRecord testData, Description description) {
101         // Track the number of iteration for this test.
102         String testName = description.getDisplayName();
103         mTestIterations.computeIfPresent(testName, (name, iterations) -> iterations + 1);
104         mTestIterations.computeIfAbsent(testName, name -> 1);
105     }
106 
107     @Override
onTestFail(DataRecord testData, Description description, Failure failure)108     public void onTestFail(DataRecord testData, Description description, Failure failure) {
109         if (mDestDir == null) {
110             return;
111         }
112         final String fileNameBase =
113                 String.format("%s.%s", description.getClassName(), description.getMethodName());
114         // Omit the iteration number for the first iteration.
115         int iteration = mTestIterations.get(description.getDisplayName());
116         final String fileName =
117                 iteration == 1
118                         ? fileNameBase
119                         : String.join("-", fileNameBase, String.valueOf(iteration));
120         // Capture the screenshot first.
121         final String pngFileName = String.format("%s-screenshot-on-failure.png", fileName);
122         File img = takeScreenshot(pngFileName);
123         if (img != null) {
124             testData.addFileMetric(String.format("%s_%s", getTag(), img.getName()), img);
125         }
126         // Capture the UI XML second.
127         if (mIncludeUiXml) {
128             File uixFile = collectUiXml(fileName);
129             if (uixFile != null) {
130                 testData.addFileMetric(
131                         String.format("%s_%s", getTag(), uixFile.getName()), uixFile);
132             }
133         }
134     }
135 
136     /** Public so that Mockito can alter its behavior. */
137     @VisibleForTesting
takeScreenshot(String fileName)138     public File takeScreenshot(String fileName) {
139         File img = new File(mDestDir, fileName);
140         try (OutputStream out = getOutputStreamViaShell(img)) {
141             screenshotToStream(out);
142             out.flush();
143             return img;
144         } catch (Exception e) {
145             Log.e(getTag(), "Unable to save screenshot", e);
146             recursiveDelete(img);
147             return null;
148         }
149     }
150 
151     /**
152      * Public so that Mockito can alter its behavior.
153      */
154     @VisibleForTesting
screenshotToStream(OutputStream out)155     public void screenshotToStream(OutputStream out) {
156         getInstrumentation().getUiAutomation()
157                 .takeScreenshot().compress(Bitmap.CompressFormat.PNG, mQuality, out);
158     }
159 
160     /** Public so that Mockito can alter its behavior. */
161     @VisibleForTesting
collectUiXml(String fileName)162     public File collectUiXml(String fileName) {
163         File uixFile = new File(mDestDir, String.format("%s.uix", fileName));
164         if (uixFile.exists()) {
165             Log.w(getTag(), String.format("File exists: %s.", uixFile.getAbsolutePath()));
166             uixFile.delete();
167         }
168         try {
169             getDevice().dumpWindowHierarchy(uixFile);
170             return uixFile;
171         } catch (IOException e) {
172             Log.e(getTag(), "Failed to collect UI XML on failure.");
173         }
174         return null;
175     }
176 
getDevice()177     private UiDevice getDevice() {
178         if (mDevice == null) {
179             mDevice = UiDevice.getInstance(getInstrumentation());
180         }
181         return mDevice;
182     }
183 }
184