1 /*
2  * Copyright (C) 2019 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 
17 package android.server.wm;
18 
19 import android.graphics.Bitmap;
20 import android.util.Log;
21 
22 import androidx.annotation.NonNull;
23 import androidx.annotation.Nullable;
24 
25 import com.android.compatibility.common.util.BitmapUtils;
26 
27 import org.junit.rules.TestRule;
28 import org.junit.runner.Description;
29 import org.junit.runners.model.Statement;
30 
31 import java.io.File;
32 import java.io.FileOutputStream;
33 import java.nio.file.Path;
34 import java.nio.file.Paths;
35 import java.util.HashMap;
36 import java.util.Map;
37 
38 /**
39  * A {@code TestRule} that allows dumping data on test failure.
40  *
41  * <p>Note: when using other {@code TestRule}s, make sure to use a {@code RuleChain} to ensure it
42  * is applied outside of other rules that can fail a test (otherwise this rule may not know that the
43  * test failed).
44  *
45  * <p>To capture the output of this rule, add the following to AndroidTest.xml:
46  * <pre>
47  *  <!-- Collect output of DumpOnFailure. -->
48  *  <metrics_collector class="com.android.tradefed.device.metric.FilePullerLogCollector">
49  *    <option name="directory-keys" value="/sdcard/DumpOnFailure" />
50  *    <option name="collect-on-run-ended-only" value="true" />
51  *  </metrics_collector>
52  * </pre>
53  * <p>And disable external storage isolation:
54  * <pre>
55  *  <uses-permission android:name="android.permission.MANAGE_EXTERNAL_STORAGE" />
56  *  <application ... android:requestLegacyExternalStorage="true" ... >
57  * </pre>
58  */
59 public class DumpOnFailure implements TestRule {
60 
61     private static final String TAG = "DumpOnFailure";
62 
63     /**
64      * Map of data to be dumped on test failure. The key must contain the name, followed by
65      * the file extension type.
66      */
67     private final Map<String, DumpItem<?>> mDumpOnFailureItems = new HashMap<>();
68 
69     @Override
apply(Statement base, Description description)70     public Statement apply(Statement base, Description description) {
71         return new Statement() {
72             @Override
73             public void evaluate() throws Throwable {
74                 onTestSetup(description);
75                 try {
76                     base.evaluate();
77                 } catch (Throwable t) {
78                     onTestFailure(description, t);
79                     throw t;
80                 } finally {
81                     onTestTeardown(description);
82                 }
83             }
84         };
85     }
86 
87     private void onTestSetup(@NonNull Description description) {
88         cleanDir(getDumpRoot(description).toFile());
89         mDumpOnFailureItems.clear();
90     }
91 
92     private void onTestTeardown(@NonNull Description description) {
93         mDumpOnFailureItems.clear();
94     }
95 
96     private void onTestFailure(@NonNull Description description, @NonNull Throwable t) {
97         final Path root = getDumpRoot(description);
98         final File rootFile = root.toFile();
99         if (!rootFile.exists() && !rootFile.mkdirs()) {
100             Log.e(TAG, "onTestFailure, unable to create file");
101             throw new RuntimeException("Unable to create " + root);
102         }
103 
104         for (var entry : mDumpOnFailureItems.entrySet()) {
105             final var fileName = getFilename(description, entry.getKey());
106             Log.i(TAG, "Dumping " + root + "/" + fileName);
107             entry.getValue().writeToFile(root.toString(), fileName);
108         }
109     }
110 
111     /**
112      * Gets the complete file name for the file to dump the data in. This includes the Test Suite,
113      * Test Method and given unique dump item name.
114      *
115      * @param description      the test description.
116      * @param nameAndExtension the unique dump item name, followed by the file extension.
117      */
118     @NonNull
119     private String getFilename(@NonNull Description description, @NonNull String nameAndExtension) {
120         return description.getTestClass().getSimpleName() + "_" + description.getMethodName()
121                 + "__" + nameAndExtension;
122     }
123 
124     @NonNull
125     private Path getDumpRoot(@NonNull Description description) {
126         return Paths.get("/sdcard/DumpOnFailure/", description.getClassName()
127                 + "_" + description.getMethodName());
128     }
129 
130     private void cleanDir(@NonNull File dir) {
131         final File[] files = dir.listFiles();
132         if (files == null) {
133             return;
134         }
135         for (File file : files) {
136             if (!file.isDirectory()) {
137                 if (!file.delete()) {
138                     throw new RuntimeException("Unable to delete " + file);
139                 }
140             }
141         }
142     }
143 
144     /**
145      * Dumps the Bitmap if the test fails.
146      *
147      * @param name   unique identifier (e.g. assertWindowShown).
148      * @param bitmap information to dump (e.g. screenshot).
149      */
150     public void dumpOnFailure(@NonNull String name, @Nullable Bitmap bitmap) {
151         if (bitmap == null) {
152             Log.i(TAG, "dumpOnFailure cannot dump null bitmap");
153             return;
154         }
155 
156         mDumpOnFailureItems.put(getNextAvailableKey(name, "png"), new BitmapItem(bitmap));
157     }
158 
159     /**
160      * Dumps the String if the test fails.
161      *
162      * @param name   unique identifier (e.g. assertWindowShown).
163      * @param string information to dump (e.g. logs).
164      */
165     public void dumpOnFailure(@NonNull String name, @Nullable String string) {
166         if (string == null) {
167             Log.i(TAG, "dumpOnFailure cannot dump null string");
168             return;
169         }
170 
171         mDumpOnFailureItems.put(getNextAvailableKey(name, "txt"), new StringItem(string));
172     }
173 
174     /**
175      * Gets the next available key in the hashmap for the given name and file extension.
176      * If the hashmap already contains an entry with the given name-extension pair, this appends
177      * the next consecutive integer that is not used for that key.
178      *
179      * @param name      the name to get the key for.
180      * @param extension the name of the file extension.
181      */
182     @NonNull
183     private String getNextAvailableKey(@NonNull String name, @NonNull String extension) {
184         if (!mDumpOnFailureItems.containsKey(name + "." + extension)) {
185             return name + "." + extension;
186         }
187 
188         int i = 1;
189         while (mDumpOnFailureItems.containsKey(name + "_" + i + "." + extension)) {
190             i++;
191         }
192         return name + "_" + i + "." + extension;
193     }
194 
195     /** Generic item containing data to be dumped on test failure. */
196     private abstract static class DumpItem<T> {
197 
198         /** The data to be dumped. */
199         @NonNull
200         final T mData;
201 
202         private DumpItem(@NonNull T data) {
203             mData = data;
204         }
205 
206         /**
207          * Writes the given data to a file created in the given directory, with the given filename.
208          *
209          * @param directoryName the name of the directory where the file should be created.
210          * @param fileName      the name of the file to be created.
211          */
212         abstract void writeToFile(@NonNull String directoryName, @NonNull String fileName);
213     }
214 
215     private static class BitmapItem extends DumpItem<Bitmap> {
216 
217         BitmapItem(@NonNull Bitmap bitmap) {
218             super(bitmap);
219         }
220 
221         @Override
222         void writeToFile(@NonNull String directoryName, @NonNull String fileName) {
223             BitmapUtils.saveBitmap(mData, directoryName, fileName);
224         }
225     }
226 
227     private static class StringItem extends DumpItem<String> {
228 
229         StringItem(@NonNull String string) {
230             super(string);
231         }
232 
233         @Override
234         void writeToFile(@NonNull String directoryName, @NonNull String fileName) {
235             Log.i(TAG, "Writing to file: " + fileName + " in directory: " + directoryName);
236 
237             final var file = new File(directoryName, fileName);
238             try (var fileStream = new FileOutputStream(file)) {
239                 fileStream.write(mData.getBytes());
240                 fileStream.flush();
241             } catch (Exception e) {
242                 Log.e(TAG, "Writing to file failed", e);
243             }
244         }
245     }
246 }
247