From 31f9d09ba9835eafd7db45c2b75a711a31e21f98 Mon Sep 17 00:00:00 2001
From: Dieter Hsu <dieterhsu@google.com>
Date: Thu, 18 Jun 2020 23:11:31 +0800
Subject: [PATCH 1/4] Run bugreport unit test without script

Replacing run.sh with fully integrated AndroidTest config to be able
to run auto test by test mapping.

Move the test class to the directory according to the package. This
fix the issue source code is referenced correctly from testing result.

Bug: 159404894
Test: atest BugreportManagerTestCases

Change-Id: Id930abc40c52441c6dd1c1e6cdc3c0f34a36bfca
Merged-In: Id930abc40c52441c6dd1c1e6cdc3c0f34a36bfca
(cherry picked from commit 7a0bac8d75d30973a6dcab361732d6adf75b83bf)
---
 core/tests/bugreports/Android.bp              |  6 ++
 core/tests/bugreports/AndroidTest.xml         | 13 ++--
 core/tests/bugreports/run.sh                  | 61 -------------------
 .../tests}/BugreportManagerTest.java          |  0
 4 files changed, 15 insertions(+), 65 deletions(-)
 delete mode 100755 core/tests/bugreports/run.sh
 rename core/tests/bugreports/src/{android/server/bugreports => com/android/os/bugreports/tests}/BugreportManagerTest.java (100%)

diff --git a/core/tests/bugreports/Android.bp b/core/tests/bugreports/Android.bp
index e9d5bb135e02c..1edd9623ed2d6 100644
--- a/core/tests/bugreports/Android.bp
+++ b/core/tests/bugreports/Android.bp
@@ -15,12 +15,18 @@
 android_test {
     name: "BugreportManagerTestCases",
     srcs: ["src/**/*.java"],
+    data: [":bugreport_artifacts"],
     libs: [
         "android.test.runner",
         "android.test.base",
     ],
     static_libs: ["androidx.test.rules", "truth-prebuilt"],
+    test_suites: ["general-tests"],
     sdk_version: "test_current",
     platform_apis: true,
 }
 
+filegroup {
+    name: "bugreport_artifacts",
+    srcs: ["config/test-sysconfig.xml"],
+}
diff --git a/core/tests/bugreports/AndroidTest.xml b/core/tests/bugreports/AndroidTest.xml
index 410ca60435838..2c2f0d69e3474 100644
--- a/core/tests/bugreports/AndroidTest.xml
+++ b/core/tests/bugreports/AndroidTest.xml
@@ -21,11 +21,16 @@
     <option name="config-descriptor:metadata" key="parameter" value="not_instant_app" />
     <option name="config-descriptor:metadata" key="parameter" value="not_multi_abi" />
 
-    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
-        <option name="cleanup-apks" value="true"/>
-        <option name="test-file-name" value="BugreportManagerTestCases.apk"/>
+    <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer" />
+    <target_preparer class="com.android.tradefed.targetprep.PushFilePreparer">
+        <option name="push-file" key="test-sysconfig.xml" value="/system/etc/sysconfig/allow-br-from-tests.xml" />
+        <option name="push-file" key="BugreportManagerTestCases.apk" value="/system/priv-app/BugreportManagerTestCases.apk" />
+        <option name="abort-on-push-failure" value="true" />
+        <option name="cleanup" value="true" />
+        <option name="remount-system" value="true" />
     </target_preparer>
-
+    <target_preparer class="com.android.tradefed.targetprep.RebootTargetPreparer" />
+
     <test class="com.android.tradefed.testtype.AndroidJUnitTest">
         <option name="package" value="com.android.os.bugreports.tests"/>
         <!-- test-timeout unit is ms, value = 30 min -->
diff --git a/core/tests/bugreports/run.sh b/core/tests/bugreports/run.sh
deleted file mode 100755
index 0103398365389..0000000000000
--- a/core/tests/bugreports/run.sh
+++ /dev/null
@@ -1,61 +0,0 @@
-#!/bin/bash
-
-# Copyright (C) 2019 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Script to run bugreport unitests
-# Must run on a rooted device.
-# Must run lunch before running the script
-# Usage: ${ANDROID_BUILD_TOP}/frameworks/base/core/tests/bugreports/run.sh
-
-# NOTE: This script replaces the framework-sysconfig.xml on your device, so use with caution.
-# It tries to replace it when done, but if the script does not finish cleanly
-# (for e.g. force stopped mid-way) your device will be left in an inconsistent state.
-# Reflashing will restore the right config.
-
-TMP_SYS_CONFIG=/var/tmp/framework-sysconfig.xml
-
-if [[ -z $ANDROID_PRODUCT_OUT ]]; then
-  echo "Please lunch before running this test."
-  exit 0
-fi
-
-# Print every command to console.
-set -x
-
-make -j BugreportManagerTestCases &&
-    adb root &&
-    adb remount &&
-    adb wait-for-device &&
-    # Save the sysconfig file in a tmp location and push the test config in
-    adb pull /system/etc/sysconfig/framework-sysconfig.xml "${TMP_SYS_CONFIG}" &&
-    adb push $ANDROID_BUILD_TOP/frameworks/base/core/tests/bugreports/config/test-sysconfig.xml /system/etc/sysconfig/framework-sysconfig.xml &&
-    # The test app needs to be a priv-app.
-    adb push $OUT/testcases/BugreportManagerTestCases/*/BugreportManagerTestCases.apk /system/priv-app ||
-    exit 1
-
-adb reboot &&
-adb wait-for-device &&
-atest BugreportManagerTest || echo "Tests FAILED!"
-
-# Restore the saved config file
-if [ -f "${TMP_SYS_CONFIG}" ]; then
-    SIZE=$(stat --printf="%s" "${TMP_SYS_CONFIG}")
-    if [ SIZE > 0 ]; then
-        adb remount &&
-        adb wait-for-device &&
-        adb push "${TMP_SYS_CONFIG}" /system/etc/sysconfig/framework-sysconfig.xml &&
-        rm "${TMP_SYS_CONFIG}"
-    fi
-fi
diff --git a/core/tests/bugreports/src/android/server/bugreports/BugreportManagerTest.java b/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
similarity index 100%
rename from core/tests/bugreports/src/android/server/bugreports/BugreportManagerTest.java
rename to core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java

From 67654483860a4623baf92dcaa9a53afc7d84a8eb Mon Sep 17 00:00:00 2001
From: Dieter Hsu <dieterhsu@google.com>
Date: Thu, 18 Jun 2020 23:11:31 +0800
Subject: [PATCH 2/4] Enable bug report tests in pre/postsubmit

Bug: 155604488
Test: atest --test-mapping --dry-run \
          services/core/java/com/android/server/os:presubmit
Test: atest --test-mapping --dry-run core/java/android/os:presubmit
Change-Id: I054d05d4fa8ff02f1c404d0883c382a2d5911f88
Merged-In: I054d05d4fa8ff02f1c404d0883c382a2d5911f88
(cherry picked from commit 848113e595fa862201277eec2d2085d231b7070e)
---
 core/java/android/os/TEST_MAPPING             | 40 +++++++++++++++++++
 packages/Shell/TEST_MAPPING                   | 31 ++++++++++++++
 .../java/com/android/server/os/TEST_MAPPING   | 36 +++++++++++++++++
 3 files changed, 107 insertions(+)
 create mode 100644 core/java/android/os/TEST_MAPPING
 create mode 100644 packages/Shell/TEST_MAPPING

diff --git a/core/java/android/os/TEST_MAPPING b/core/java/android/os/TEST_MAPPING
new file mode 100644
index 0000000000000..f4645caee5113
--- /dev/null
+++ b/core/java/android/os/TEST_MAPPING
@@ -0,0 +1,40 @@
+{
+  "presubmit": [
+    // TODO(159590499) add BugreportManagerTestCases
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "CtsBugreportTestCases",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        }
+      ]
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "ShellTests",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        },
+        {
+          "exclude-annotation": "androidx.test.filters.FlakyTest"
+        }
+      ]
+    }
+  ],
+  "postsubmit": [
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "BugreportManagerTestCases"
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "CtsBugreportTestCases"
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "ShellTests"
+    }
+  ]
+}
diff --git a/packages/Shell/TEST_MAPPING b/packages/Shell/TEST_MAPPING
new file mode 100644
index 0000000000000..a149b5cda3eac
--- /dev/null
+++ b/packages/Shell/TEST_MAPPING
@@ -0,0 +1,31 @@
+{
+  "presubmit": [
+    {
+      "name": "CtsBugreportTestCases",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        }
+      ]
+    },
+    {
+      "name": "ShellTests",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        },
+        {
+          "exclude-annotation": "androidx.test.filters.FlakyTest"
+        }
+      ]
+    }
+  ],
+  "postsubmit": [
+    {
+      "name": "CtsBugreportTestCases"
+    },
+    {
+      "name": "ShellTests"
+    }
+  ]
+}
diff --git a/services/core/java/com/android/server/os/TEST_MAPPING b/services/core/java/com/android/server/os/TEST_MAPPING
index 502f1e852e086..a837fb4184ba8 100644
--- a/services/core/java/com/android/server/os/TEST_MAPPING
+++ b/services/core/java/com/android/server/os/TEST_MAPPING
@@ -1,7 +1,43 @@
 {
   "presubmit": [
+    // TODO(159590499) add BugreportManagerTestCases
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "CtsBugreportTestCases",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        }
+      ]
+    },
     {
       "name": "CtsUsbTests"
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "ShellTests",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        },
+        {
+          "exclude-annotation": "androidx.test.filters.FlakyTest"
+        }
+      ]
+    }
+  ],
+  "postsubmit": [
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "BugreportManagerTestCases"
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "CtsBugreportTestCases"
+    },
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "ShellTests"
     }
   ]
 }

From 59a1d2b4aa0ade79b655c5628f78a9f8ca0d07c4 Mon Sep 17 00:00:00 2001
From: Dieter Hsu <dieterhsu@google.com>
Date: Thu, 2 Jul 2020 01:09:47 +0800
Subject: [PATCH 3/4] Add BugreportManagerTestCases to bugreport presubmit

Bug: 159590499
Test: atest --test-mapping --dry-run \
                 services/core/java/com/android/server/os:presubmit
Test: atest --test-mapping --dry-run core/java/android/os:presubmit

Change-Id: Ib6ec81570d5000ae3c7cb76da0e8e0eed98c425e
Merged-In: Ib6ec81570d5000ae3c7cb76da0e8e0eed98c425e
(cherry picked from commit f7f4f62850d7136a1b0bfe33c0e899736992d4f7)
---
 core/java/android/os/TEST_MAPPING                     | 10 +++++++++-
 services/core/java/com/android/server/os/TEST_MAPPING | 10 +++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/core/java/android/os/TEST_MAPPING b/core/java/android/os/TEST_MAPPING
index f4645caee5113..39c196d1a06c6 100644
--- a/core/java/android/os/TEST_MAPPING
+++ b/core/java/android/os/TEST_MAPPING
@@ -1,6 +1,14 @@
 {
   "presubmit": [
-    // TODO(159590499) add BugreportManagerTestCases
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "BugreportManagerTestCases",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        }
+      ]
+    },
     {
       "file_patterns": ["Bugreport[^/]*\\.java"],
       "name": "CtsBugreportTestCases",
diff --git a/services/core/java/com/android/server/os/TEST_MAPPING b/services/core/java/com/android/server/os/TEST_MAPPING
index a837fb4184ba8..d937af1e3fdb5 100644
--- a/services/core/java/com/android/server/os/TEST_MAPPING
+++ b/services/core/java/com/android/server/os/TEST_MAPPING
@@ -1,6 +1,14 @@
 {
   "presubmit": [
-    // TODO(159590499) add BugreportManagerTestCases
+    {
+      "file_patterns": ["Bugreport[^/]*\\.java"],
+      "name": "BugreportManagerTestCases",
+      "options": [
+        {
+          "exclude-annotation": "androidx.test.filters.LargeTest"
+        }
+      ]
+    },
     {
       "file_patterns": ["Bugreport[^/]*\\.java"],
       "name": "CtsBugreportTestCases",

From 8f74667510edc35457cf367b971269a59b87d35b Mon Sep 17 00:00:00 2001
From: Dieter Hsu <dieterhsu@google.com>
Date: Mon, 29 Jun 2020 13:18:34 +0800
Subject: [PATCH 4/4] Fix consent timed out to sharing bugreport data

Also mark the fixed tests as large test to be ignored when presubmit
running.

Bug: 159590499
Test: atest BugreportManagerTest
Change-Id: Ie926eefacb659cce4c54f42b0835308da5a71bba
Merged-In: Ie926eefacb659cce4c54f42b0835308da5a71bba
(cherry picked from commit 572ef12b3dc0e270c5837569938eca4ba8bbf71a)
---
 core/tests/bugreports/Android.bp              |   6 +-
 .../tests/BugreportManagerTest.java           | 128 ++++++++++++++++--
 2 files changed, 119 insertions(+), 15 deletions(-)

diff --git a/core/tests/bugreports/Android.bp b/core/tests/bugreports/Android.bp
index 1edd9623ed2d6..e42b4b4ae2bf1 100644
--- a/core/tests/bugreports/Android.bp
+++ b/core/tests/bugreports/Android.bp
@@ -20,7 +20,11 @@ android_test {
         "android.test.runner",
         "android.test.base",
     ],
-    static_libs: ["androidx.test.rules", "truth-prebuilt"],
+    static_libs: [
+        "androidx.test.rules",
+        "androidx.test.uiautomator_uiautomator",
+        "truth-prebuilt",
+    ],
     test_suites: ["general-tests"],
     sdk_version: "test_current",
     platform_apis: true,
diff --git a/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java b/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
index 153337727e967..9246a2335050e 100644
--- a/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
+++ b/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
@@ -18,6 +18,8 @@ package com.android.os.bugreports.tests;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
 import android.Manifest;
@@ -25,17 +27,27 @@ import android.content.Context;
 import android.os.BugreportManager;
 import android.os.BugreportManager.BugreportCallback;
 import android.os.BugreportParams;
+import android.os.FileUtils;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.ParcelFileDescriptor;
+import android.os.StrictMode;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
 import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.LargeTest;
+import androidx.test.uiautomator.By;
+import androidx.test.uiautomator.BySelector;
+import androidx.test.uiautomator.UiDevice;
+import androidx.test.uiautomator.UiObject2;
+import androidx.test.uiautomator.Until;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExternalResource;
 import org.junit.rules.TestName;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -51,9 +63,11 @@ import java.util.concurrent.TimeUnit;
 @RunWith(JUnit4.class)
 public class BugreportManagerTest {
     @Rule public TestName name = new TestName();
+    @Rule public ExtendedStrictModeVmPolicy mTemporaryVmPolicy = new ExtendedStrictModeVmPolicy();
 
     private static final String TAG = "BugreportManagerTest";
     private static final long BUGREPORT_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(10);
+    private static final long UIAUTOMATOR_TIMEOUT_MS = TimeUnit.SECONDS.toMillis(10);
 
     private Handler mHandler;
     private Executor mExecutor;
@@ -86,6 +100,8 @@ public class BugreportManagerTest {
     @After
     public void teardown() throws Exception {
         dropPermissions();
+        FileUtils.closeQuietly(mBugreportFd);
+        FileUtils.closeQuietly(mScreenshotFd);
     }
 
 
@@ -95,47 +111,45 @@ public class BugreportManagerTest {
         // wifi bugreport does not take screenshot
         mBrm.startBugreport(mBugreportFd, null /*screenshotFd = null*/, wifi(),
                 mExecutor, callback);
+        shareConsentDialog(ConsentReply.ALLOW);
         waitTillDoneOrTimeout(callback);
 
         assertThat(callback.isDone()).isTrue();
         // Wifi bugreports should not receive any progress.
         assertThat(callback.hasReceivedProgress()).isFalse();
-        // TODO: Because of b/130234145, consent dialog is not shown; so we get a timeout error.
-        // When the bug is fixed, accept consent via UIAutomator and verify contents
-        // of mBugreportFd.
-        assertThat(callback.getErrorCode()).isEqualTo(
-                BugreportCallback.BUGREPORT_ERROR_USER_CONSENT_TIMED_OUT);
+        assertThat(mBugreportFile.length()).isGreaterThan(0L);
         assertFdsAreClosed(mBugreportFd);
     }
 
+    @LargeTest
     @Test
     public void normalFlow_interactive() throws Exception {
         BugreportCallbackImpl callback = new BugreportCallbackImpl();
         // interactive bugreport does not take screenshot
         mBrm.startBugreport(mBugreportFd, null /*screenshotFd = null*/, interactive(),
                 mExecutor, callback);
-
+        shareConsentDialog(ConsentReply.ALLOW);
         waitTillDoneOrTimeout(callback);
+
         assertThat(callback.isDone()).isTrue();
         // Interactive bugreports show progress updates.
         assertThat(callback.hasReceivedProgress()).isTrue();
-        assertThat(callback.getErrorCode()).isEqualTo(
-                BugreportCallback.BUGREPORT_ERROR_USER_CONSENT_TIMED_OUT);
+        assertThat(mBugreportFile.length()).isGreaterThan(0L);
         assertFdsAreClosed(mBugreportFd);
     }
 
+    @LargeTest
     @Test
     public void normalFlow_full() throws Exception {
         BugreportCallbackImpl callback = new BugreportCallbackImpl();
         mBrm.startBugreport(mBugreportFd, mScreenshotFd, full(), mExecutor, callback);
-
+        shareConsentDialog(ConsentReply.ALLOW);
         waitTillDoneOrTimeout(callback);
+
         assertThat(callback.isDone()).isTrue();
-        assertThat(callback.getErrorCode()).isEqualTo(
-                BugreportCallback.BUGREPORT_ERROR_USER_CONSENT_TIMED_OUT);
-        // bugreport and screenshot files should be empty when user consent timed out.
-        assertThat(mBugreportFile.length()).isEqualTo(0);
-        assertThat(mScreenshotFile.length()).isEqualTo(0);
+        // bugreport and screenshot files shouldn't be empty when user consents.
+        assertThat(mBugreportFile.length()).isGreaterThan(0L);
+        assertThat(mScreenshotFile.length()).isGreaterThan(0L);
         assertFdsAreClosed(mBugreportFd, mScreenshotFd);
     }
 
@@ -144,6 +158,8 @@ public class BugreportManagerTest {
         // Start bugreport #1
         BugreportCallbackImpl callback = new BugreportCallbackImpl();
         mBrm.startBugreport(mBugreportFd, mScreenshotFd, wifi(), mExecutor, callback);
+        // TODO(b/162389762) Make sure the wait time is reasonable
+        shareConsentDialog(ConsentReply.ALLOW);
 
         // Before #1 is done, try to start #2.
         assertThat(callback.isDone()).isFalse();
@@ -375,4 +391,88 @@ public class BugreportManagerTest {
     private static BugreportParams full() {
         return new BugreportParams(BugreportParams.BUGREPORT_MODE_FULL);
     }
+
+    /* Allow/deny the consent dialog to sharing bugreport data or check existence only. */
+    private enum ConsentReply {
+        ALLOW,
+        DENY,
+        TIMEOUT
+    }
+
+    /*
+     * Ensure the consent dialog is shown and take action according to <code>consentReply<code/>.
+     * It will fail if the dialog is not shown when <code>ignoreNotFound<code/> is false.
+     */
+    private void shareConsentDialog(@NonNull ConsentReply consentReply) throws Exception {
+        mTemporaryVmPolicy.permitIncorrectContextUse();
+        final UiDevice device = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
+
+        // Unlock before finding/clicking an object.
+        device.wakeUp();
+        device.executeShellCommand("wm dismiss-keyguard");
+
+        final BySelector consentTitleObj = By.res("android", "alertTitle");
+        if (!device.wait(Until.hasObject(consentTitleObj), UIAUTOMATOR_TIMEOUT_MS)) {
+            fail("The consent dialog is not found");
+        }
+        if (consentReply.equals(ConsentReply.TIMEOUT)) {
+            return;
+        }
+        final BySelector selector;
+        if (consentReply.equals(ConsentReply.ALLOW)) {
+            selector = By.res("android", "button1");
+            Log.d(TAG, "Allow the consent dialog");
+        } else { // ConsentReply.DENY
+            selector = By.res("android", "button2");
+            Log.d(TAG, "Deny the consent dialog");
+        }
+        final UiObject2 btnObj = device.findObject(selector);
+        assertNotNull("The button of consent dialog is not found", btnObj);
+        btnObj.click();
+
+        Log.d(TAG, "Wait for the dialog to be dismissed");
+        assertTrue(device.wait(Until.gone(consentTitleObj), UIAUTOMATOR_TIMEOUT_MS));
+    }
+
+    /**
+     * A rule to change strict mode vm policy temporarily till test method finished.
+     *
+     * To permit the non-visual context usage in tests while taking bugreports need user consent,
+     * or UiAutomator/BugreportManager.DumpstateListener would run into error.
+     * UiDevice#findObject creates UiObject2, its Gesture object and ViewConfiguration and
+     * UiObject2#click need to know bounds. Both of them access to WindowManager internally without
+     * visual context comes from InstrumentationRegistry and violate the policy.
+     * Also <code>DumpstateListener<code/> violate the policy when onScreenshotTaken is called.
+     *
+     * TODO(b/161201609) Remove this class once violations fixed.
+     */
+    static class ExtendedStrictModeVmPolicy extends ExternalResource {
+        private boolean mWasVmPolicyChanged = false;
+        private StrictMode.VmPolicy mOldVmPolicy;
+
+        @Override
+        protected void after() {
+            restoreVmPolicyIfNeeded();
+        }
+
+        public void permitIncorrectContextUse() {
+            // Allow to call multiple times without losing old policy.
+            if (mOldVmPolicy == null) {
+                mOldVmPolicy = StrictMode.getVmPolicy();
+            }
+            StrictMode.setVmPolicy(new StrictMode.VmPolicy.Builder()
+                    .detectAll()
+                    .permitIncorrectContextUse()
+                    .penaltyLog()
+                    .build());
+            mWasVmPolicyChanged = true;
+        }
+
+        private void restoreVmPolicyIfNeeded() {
+            if (mWasVmPolicyChanged && mOldVmPolicy != null) {
+                StrictMode.setVmPolicy(mOldVmPolicy);
+                mOldVmPolicy = null;
+            }
+        }
+    }
 }