Home | History | Annotate | Download | only in collectors 
      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.support.annotation.VisibleForTesting;
     22 import android.util.Log;
     23 
     24 import org.junit.runner.Description;
     25 
     26 import java.io.BufferedOutputStream;
     27 import java.io.ByteArrayOutputStream;
     28 import java.io.File;
     29 import java.io.FileOutputStream;
     30 import java.io.OutputStream;
     31 
     32 /**
     33  * A {@link BaseMetricListener} that captures screenshots at the end of each test case.
     34  *
     35  * This class needs external storage permission. See {@link BaseMetricListener} how to grant
     36  * external storage permission, especially at install time.
     37  *
     38  * Options:
     39  * -e screenshot-quality [0-100]: set screenshot image quality. Default is 75.
     40  * -e screenshot-format [byte|file:optional/path/to/dir]: set screenshot format. Default is file.
     41  * Choose "byte" to save as a byte array in result bundle.
     42  * Choose "file" to save as an image file. Append ":path/to/dir" behind "file" to specify directory
     43  * to save screenshots, relative to /sdcard/. e.g. "-e screenshot-format file:tmp/sc" will save
     44  * screenshots to /sdcard/tmp/sc/ directory.
     45  *
     46  * Do NOT throw exception anywhere in this class. We don't want to halt the test when metrics
     47  * collection fails.
     48  */
     49 @OptionClass(alias = "screenshot-collector")
     50 public class ScreenshotListener extends BaseMetricListener {
     51 
     52     public static final String DEFAULT_DIR = "run_listeners/screenshots";
     53     public static final String KEY_QUALITY = "screenshot-quality"; // 0 to 100
     54     public static final String KEY_FORMAT = "screenshot-format"; // byte or file
     55     public static final int DEFAULT_QUALITY = 75;
     56     public static final String OPTION_BYTE = "byte";
     57     private int mQuality = DEFAULT_QUALITY;
     58     private boolean mToFile;
     59 
     60     private File mDestDir;
     61 
     62     public ScreenshotListener() {
     63         super();
     64     }
     65 
     66     /**
     67      * Constructor to simulate receiving the instrumentation arguments. Should not be used except
     68      * for testing.
     69      */
     70     @VisibleForTesting
     71     ScreenshotListener(Bundle args) {
     72         super(args);
     73     }
     74 
     75     @Override
     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         mToFile = !OPTION_BYTE.equals(args.getString(KEY_FORMAT));
     91         if (mToFile) {
     92             String dir = DEFAULT_DIR;
     93             if (args.containsKey(KEY_FORMAT)) {
     94                 String[] argsArray = args.getString(KEY_FORMAT).split(":", 2);
     95                 if (argsArray.length > 1) {
     96                     dir = argsArray[1];
     97                 }
     98             }
     99             mDestDir = createAndEmptyDirectory(dir);
    100         }
    101     }
    102 
    103     @Override
    104     public void onTestEnd(DataRecord testData, Description description) {
    105         if (mToFile) {
    106             if (mDestDir == null) {
    107                 return;
    108             }
    109             final String fileName = String.format("%s.%s.png", description.getClassName(),
    110                     description.getMethodName());
    111             File img = takeScreenshot(fileName);
    112             if (img != null) {
    113                 testData.addFileMetric(String.format("%s_%s", getTag(), img.getName()), img);
    114             }
    115         } else {
    116             ByteArrayOutputStream out = new ByteArrayOutputStream();
    117             screenshotToStream(out);
    118             testData.addBinaryMetric(String.format("%s_%s.%s.bytes", getTag(),
    119                     description.getClassName(), description.getMethodName()), out.toByteArray());
    120         }
    121     }
    122 
    123     /**
    124      * Public so that Mockito can alter its behavior.
    125      */
    126     @VisibleForTesting
    127     public File takeScreenshot(String fileName){
    128         File img = new File(mDestDir, fileName);
    129         if (img.exists()) {
    130             Log.w(getTag(), String.format("File exists: %s", img.getAbsolutePath()));
    131             img.delete();
    132         }
    133         try (
    134                 OutputStream out = new BufferedOutputStream(new FileOutputStream(img))
    135         ){
    136             screenshotToStream(out);
    137             out.flush();
    138             return img;
    139         } catch (Exception e) {
    140             Log.e(getTag(), "Unable to save screenshot", e);
    141             img.delete();
    142             return null;
    143         }
    144     }
    145 
    146     /**
    147      * Public so that Mockito can alter its behavior.
    148      */
    149     @VisibleForTesting
    150     public void screenshotToStream(OutputStream out) {
    151         getInstrumentation().getUiAutomation()
    152                 .takeScreenshot().compress(Bitmap.CompressFormat.PNG, mQuality, out);
    153     }
    154 }
    155