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.BufferedOutputStream;
29 import java.io.File;
30 import java.io.FileOutputStream;
31 import java.io.IOException;
32 import java.io.OutputStream;
33 import java.util.Map;
34 import java.util.HashMap;
35 
36 
37 /**
38  * A {@link BaseMetricListener} that captures screenshots when a test fails.
39  *
40  * <p>This class needs external storage permission. See {@link BaseMetricListener} how to grant
41  * external storage permission, especially at install time.
42  *
43  * <p>Options: -e screenshot-quality [0-100]: set screenshot image quality. Default is 75. -e
44  * include-ui-xml [true, false]: include the UI XML on failure too, if true.
45  */
46 @OptionClass(alias = "screenshot-failure-collector")
47 public class ScreenshotOnFailureCollector extends BaseMetricListener {
48 
49     public static final String DEFAULT_DIR = "run_listeners/screenshots";
50     public static final String KEY_INCLUDE_XML = "include-ui-xml";
51     public static final String KEY_QUALITY = "screenshot-quality";
52     public static final int DEFAULT_QUALITY = 75;
53     private boolean mIncludeUiXml = false;
54     private int mQuality = DEFAULT_QUALITY;
55 
56     private File mDestDir;
57     private UiDevice mDevice;
58 
59     // Tracks the test iterations to ensure that each failure gets unique filenames.
60     // Key: test description; value: number of iterations.
61     private Map<String, Integer> mTestIterations = new HashMap<String, Integer>();
62 
ScreenshotOnFailureCollector()63     public ScreenshotOnFailureCollector() {
64         super();
65     }
66 
67     /**
68      * Constructor to simulate receiving the instrumentation arguments. Should not be used except
69      * for testing.
70      */
71     @VisibleForTesting
ScreenshotOnFailureCollector(Bundle args)72     ScreenshotOnFailureCollector(Bundle args) {
73         super(args);
74     }
75 
76     @Override
onTestRunStart(DataRecord runData, Description description)77     public void onTestRunStart(DataRecord runData, Description description) {
78         Bundle args = getArgsBundle();
79         if (args.containsKey(KEY_QUALITY)) {
80             try {
81                 int quality = Integer.parseInt(args.getString(KEY_QUALITY));
82                 if (quality >= 0 && quality <= 100) {
83                     mQuality = quality;
84                 } else {
85                     Log.e(getTag(), String.format("Invalid screenshot quality: %d.", quality));
86                 }
87             } catch (Exception e) {
88                 Log.e(getTag(), "Failed to parse screenshot quality", e);
89             }
90         }
91 
92         if (args.containsKey(KEY_INCLUDE_XML)) {
93             mIncludeUiXml = Boolean.parseBoolean(args.getString(KEY_INCLUDE_XML));
94         }
95 
96         String dir = DEFAULT_DIR;
97         mDestDir = createAndEmptyDirectory(dir);
98     }
99 
100     @Override
onTestStart(DataRecord testData, Description description)101     public void onTestStart(DataRecord testData, Description description) {
102         // Track the number of iteration for this test.
103         String testName = description.getDisplayName();
104         mTestIterations.computeIfPresent(testName, (name, iterations) -> iterations + 1);
105         mTestIterations.computeIfAbsent(testName, name -> 1);
106     }
107 
108     @Override
onTestFail(DataRecord testData, Description description, Failure failure)109     public void onTestFail(DataRecord testData, Description description, Failure failure) {
110         if (mDestDir == null) {
111             return;
112         }
113         final String fileNameBase =
114                 String.format("%s.%s", description.getClassName(), description.getMethodName());
115         // Omit the iteration number for the first iteration.
116         int iteration = mTestIterations.get(description.getDisplayName());
117         final String fileName =
118                 iteration == 1
119                         ? fileNameBase
120                         : String.join("-", fileNameBase, String.valueOf(iteration));
121         // Capture the screenshot first.
122         final String pngFileName = String.format("%s-screenshot-on-failure.png", fileName);
123         File img = takeScreenshot(pngFileName);
124         if (img != null) {
125             testData.addFileMetric(String.format("%s_%s", getTag(), img.getName()), img);
126         }
127         // Capture the UI XML second.
128         if (mIncludeUiXml) {
129             File uixFile = collectUiXml(fileName);
130             if (uixFile != null) {
131                 testData.addFileMetric(
132                         String.format("%s_%s", getTag(), uixFile.getName()), uixFile);
133             }
134         }
135     }
136 
137     /** Public so that Mockito can alter its behavior. */
138     @VisibleForTesting
takeScreenshot(String fileName)139     public File takeScreenshot(String fileName) {
140         File img = new File(mDestDir, fileName);
141         if (img.exists()) {
142             Log.w(getTag(), String.format("File exists: %s", img.getAbsolutePath()));
143             img.delete();
144         }
145         try (
146                 OutputStream out = new BufferedOutputStream(new FileOutputStream(img))
147         ){
148             screenshotToStream(out);
149             out.flush();
150             return img;
151         } catch (Exception e) {
152             Log.e(getTag(), "Unable to save screenshot", e);
153             img.delete();
154             return null;
155         }
156     }
157 
158     /**
159      * Public so that Mockito can alter its behavior.
160      */
161     @VisibleForTesting
screenshotToStream(OutputStream out)162     public void screenshotToStream(OutputStream out) {
163         getInstrumentation().getUiAutomation()
164                 .takeScreenshot().compress(Bitmap.CompressFormat.PNG, mQuality, out);
165     }
166 
167     /** Public so that Mockito can alter its behavior. */
168     @VisibleForTesting
collectUiXml(String fileName)169     public File collectUiXml(String fileName) {
170         File uixFile = new File(mDestDir, String.format("%s.uix", fileName));
171         if (uixFile.exists()) {
172             Log.w(getTag(), String.format("File exists: %s.", uixFile.getAbsolutePath()));
173             uixFile.delete();
174         }
175         try {
176             getDevice().dumpWindowHierarchy(uixFile);
177             return uixFile;
178         } catch (IOException e) {
179             Log.e(getTag(), "Failed to collect UI XML on failure.");
180         }
181         return null;
182     }
183 
getDevice()184     private UiDevice getDevice() {
185         if (mDevice == null) {
186             mDevice = UiDevice.getInstance(getInstrumentation());
187         }
188         return mDevice;
189     }
190 }
191