1 /*
2  * Copyright (C) 2018 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 com.android.tradefed.device.metric;
17 
18 import com.android.tradefed.device.DeviceNotAvailableException;
19 import com.android.tradefed.device.ITestDevice;
20 import com.android.tradefed.device.ITestDevice.RecoveryMode;
21 import com.android.tradefed.device.TestDeviceState;
22 import com.android.tradefed.log.LogUtil.CLog;
23 import com.android.tradefed.result.InputStreamSource;
24 import com.android.tradefed.result.LogDataType;
25 import com.android.tradefed.result.TestDescription;
26 
27 /** Collector that will capture and log a screenshot when a test case fails. */
28 public class ScreenshotOnFailureCollector extends BaseDeviceMetricCollector {
29 
30     private static final String NAME_FORMAT = "%s-%s-screenshot-on-failure";
31     private static final int THROTTLE_LIMIT_PER_RUN = 10;
32 
33     private int mCurrentCount = 0;
34     private boolean mFirstThrottle = true;
35 
36     @Override
onTestRunStart(DeviceMetricData runData)37     public void onTestRunStart(DeviceMetricData runData) {
38         mCurrentCount = 0;
39         mFirstThrottle = true;
40     }
41 
42     @Override
onTestFail(DeviceMetricData testData, TestDescription test)43     public void onTestFail(DeviceMetricData testData, TestDescription test)
44             throws DeviceNotAvailableException {
45         if (mCurrentCount > THROTTLE_LIMIT_PER_RUN) {
46             if (mFirstThrottle) {
47                 CLog.w("Throttle capture of screenshot-on-failure due to too many failures.");
48                 mFirstThrottle = false;
49             }
50             return;
51         }
52         for (ITestDevice device : getRealDevices()) {
53             if (!shouldCollect(device)) {
54                 continue;
55             }
56             RecoveryMode mode = device.getRecoveryMode();
57             device.setRecoveryMode(RecoveryMode.NONE);
58             try (InputStreamSource screenSource = device.getScreenshot()) {
59                 CLog.d("Captured screenshot-on-failure.");
60                 super.testLog(
61                         String.format(NAME_FORMAT, test.toString(), device.getSerialNumber()),
62                         LogDataType.PNG,
63                         screenSource);
64             } finally {
65                 device.setRecoveryMode(mode);
66             }
67         }
68         mCurrentCount++;
69     }
70 
shouldCollect(ITestDevice device)71     private boolean shouldCollect(ITestDevice device) {
72         TestDeviceState state = device.getDeviceState();
73         if (!TestDeviceState.ONLINE.equals(state)) {
74             CLog.d("Skip ScreenshotOnFailureCollector device is in state '%s'", state);
75             return false;
76         }
77         return true;
78     }
79 }
80