Test migration for setScreenCaptureDisabled

Migrated com.android.cts.devicepolicy.DeviceAndProfileOwnerTest#testScreenCaptureDisabled

Bug: 197766622
Test: atest android.devicepolicy.cts.ScreenCaptureDisabledTest
Change-Id: If75f4d70ea331fbc6d33be9b3990b91a12197353
Merged-In: If75f4d70ea331fbc6d33be9b3990b91a12197353
(cherry picked from commit d7e31245820ee86f49afa885098d7ab422201a1f)
diff --git a/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/policies/ScreenCaptureDisabled.java b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/policies/ScreenCaptureDisabled.java
new file mode 100644
index 0000000..e165663
--- /dev/null
+++ b/common/device-side/bedstead/harrier/src/main/java/com/android/bedstead/harrier/policies/ScreenCaptureDisabled.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+package com.android.bedstead.harrier.policies;
+
+import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_DEVICE_OWNER;
+import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_PARENT_INSTANCE_OF_PROFILE_OWNER;
+import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIED_BY_PROFILE_OWNER;
+import static com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy.APPLIES_TO_OWN_USER;
+
+import android.app.admin.DevicePolicyManager;
+import android.content.ComponentName;
+
+import com.android.bedstead.harrier.annotations.enterprise.EnterprisePolicy;
+
+/**
+ * Policy for disabling screen capture.
+ *
+ * <p>Users of this policy are
+ * {@link DevicePolicyManager#setScreenCaptureDisabled(ComponentName, boolean)}.
+ */
+@EnterprisePolicy(dpc = APPLIED_BY_DEVICE_OWNER | APPLIED_BY_PROFILE_OWNER
+        | APPLIED_BY_PARENT_INSTANCE_OF_PROFILE_OWNER | APPLIES_TO_OWN_USER)
+public final class ScreenCaptureDisabled {
+}
diff --git a/common/device-side/util-axt/src/com/android/compatibility/common/util/PollingCheck.java b/common/device-side/util-axt/src/com/android/compatibility/common/util/PollingCheck.java
index 00eda91..e42eccc 100644
--- a/common/device-side/util-axt/src/com/android/compatibility/common/util/PollingCheck.java
+++ b/common/device-side/util-axt/src/com/android/compatibility/common/util/PollingCheck.java
@@ -16,10 +16,12 @@
 
 package com.android.compatibility.common.util;
 
+import junit.framework.Assert;
+
 import java.util.concurrent.Callable;
 import java.util.function.BooleanSupplier;
-
-import junit.framework.Assert;
+import java.util.function.Function;
+import java.util.function.Supplier;
 
 public abstract class PollingCheck {
     private static final long TIME_SLICE = 50;
@@ -71,6 +73,30 @@
         Assert.assertTrue(mErrorMessage, check());
     }
 
+    public <E> E runWaitAndReturnResult(Supplier<E> supplier, Function<E, Boolean> condition) {
+        E output = supplier.get();
+        if (condition.apply(output)) {
+            return output;
+        }
+        long timeout = mTimeout;
+        while (timeout > 0) {
+            try {
+                Thread.sleep(TIME_SLICE);
+            } catch (InterruptedException e) {
+                Assert.fail("unexpected InterruptedException");
+            }
+
+            output = supplier.get();
+            if (condition.apply(output)) {
+                return output;
+            }
+
+            timeout -= TIME_SLICE;
+        }
+
+        return output;
+    }
+
     public static void check(CharSequence message, long timeout, Callable<Boolean> condition)
             throws Exception {
         while (timeout > 0) {
@@ -111,4 +137,16 @@
             }
         }.run();
     }
+
+    public static <E> E waitFor(long timeout, Supplier<E> supplier,
+            Function<E, Boolean> condition) {
+        return new PollingCheck(timeout) {
+            // Not used in the corresponding runWaitAndReturnResult function, Added just to provide
+            // default implementation of abstract check function.
+            @Override
+            protected boolean check() {
+                return true;
+            }
+        }.runWaitAndReturnResult(supplier, condition);
+    }
 }
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
index 0b59df3..38844ea 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/DeviceAndProfileOwnerTest.java
@@ -682,26 +682,6 @@
     }
 
     @Test
-    public void testScreenCaptureDisabled() throws Exception {
-        assertMetricsLogged(getDevice(), () -> {
-            // We need to ensure that the policy is deactivated for the device owner case, so making
-            // sure the second test is run even if the first one fails
-            try {
-                setScreenCaptureDisabled(mUserId, true);
-            } finally {
-                setScreenCaptureDisabled(mUserId, false);
-            }
-        }, new DevicePolicyEventWrapper.Builder(EventId.SET_SCREEN_CAPTURE_DISABLED_VALUE)
-                    .setAdminPackageName(DEVICE_ADMIN_PKG)
-                    .setBoolean(true)
-                    .build(),
-            new DevicePolicyEventWrapper.Builder(EventId.SET_SCREEN_CAPTURE_DISABLED_VALUE)
-                    .setAdminPackageName(DEVICE_ADMIN_PKG)
-                    .setBoolean(false)
-                    .build());
-    }
-
-    @Test
     public void testScreenCaptureDisabled_assist() throws Exception {
         try {
             // Install and enable assistant, notice that profile can't have assistant.
@@ -2026,21 +2006,6 @@
         startActivityAsUser(userId, TEST_APP_PKG, TEST_APP_PKG + ".SimpleActivity");
     }
 
-    protected void setScreenCaptureDisabled(int userId, boolean disabled) throws Exception {
-        String testMethodName = disabled
-                ? "testSetScreenCaptureDisabled_true"
-                : "testSetScreenCaptureDisabled_false";
-        executeDeviceTestMethod(".ScreenCaptureDisabledTest", testMethodName);
-
-        testMethodName = disabled
-                ? "testScreenCaptureImpossible"
-                : "testScreenCapturePossible";
-
-        startSimpleActivityAsUser(userId);
-        executeDeviceTestMethod(".ScreenCaptureDisabledTest", testMethodName);
-        forceStopPackageForUser(TEST_APP_PKG, userId);
-    }
-
     protected void setScreenCaptureDisabled_assist(int userId, boolean disabled) throws Exception {
         // Set the policy.
         String testMethodName = disabled
diff --git a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedManagedProfileOwnerTest.java b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedManagedProfileOwnerTest.java
index 07f7efd..df2e7c9 100644
--- a/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedManagedProfileOwnerTest.java
+++ b/hostsidetests/devicepolicy/src/com/android/cts/devicepolicy/MixedManagedProfileOwnerTest.java
@@ -77,22 +77,6 @@
 
     // Most tests for this class are defined in DeviceAndProfileOwnerTest
 
-    /**
-     * Verify that screenshots are still possible for activities in the primary user when the policy
-     * is set on the profile owner.
-     */
-    @LargeTest
-    @Test
-    public void testScreenCaptureDisabled_allowedPrimaryUser() throws Exception {
-        // disable screen capture in profile
-        setScreenCaptureDisabled(mUserId, true);
-
-        // start the ScreenCaptureDisabledActivity in the parent
-        installAppAsUser(DEVICE_ADMIN_APK, mParentUserId);
-        startSimpleActivityAsUser(mParentUserId);
-        executeDeviceTestMethod(".ScreenCaptureDisabledTest", "testScreenCapturePossible");
-    }
-
     @FlakyTest
     @Test
     public void testScreenCaptureDisabled_assist_allowedPrimaryUser() throws Exception {
diff --git a/tests/devicepolicy/src/android/devicepolicy/cts/ScreenCaptureDisabledTest.java b/tests/devicepolicy/src/android/devicepolicy/cts/ScreenCaptureDisabledTest.java
new file mode 100644
index 0000000..c50621e
--- /dev/null
+++ b/tests/devicepolicy/src/android/devicepolicy/cts/ScreenCaptureDisabledTest.java
@@ -0,0 +1,208 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+
+package android.devicepolicy.cts;
+
+
+import static com.android.bedstead.metricsrecorder.truth.MetricQueryBuilderSubject.assertThat;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.UiAutomation;
+import android.app.admin.DevicePolicyManager;
+import android.app.admin.RemoteDevicePolicyManager;
+import android.content.ComponentName;
+import android.graphics.Bitmap;
+import android.stats.devicepolicy.EventId;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.bedstead.harrier.BedsteadJUnit4;
+import com.android.bedstead.harrier.DeviceState;
+import com.android.bedstead.harrier.annotations.Postsubmit;
+import com.android.bedstead.harrier.annotations.enterprise.CanSetPolicyTest;
+import com.android.bedstead.harrier.annotations.enterprise.NegativePolicyTest;
+import com.android.bedstead.harrier.annotations.enterprise.PositivePolicyTest;
+import com.android.bedstead.harrier.policies.ScreenCaptureDisabled;
+import com.android.bedstead.metricsrecorder.EnterpriseMetricsRecorder;
+import com.android.bedstead.nene.TestApis;
+import com.android.bedstead.testapp.TestApp;
+import com.android.bedstead.testapp.TestAppInstanceReference;
+import com.android.bedstead.testapp.TestAppProvider;
+import com.android.compatibility.common.util.PollingCheck;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Objects;
+
+
+@RunWith(BedsteadJUnit4.class)
+public class ScreenCaptureDisabledTest {
+
+    @ClassRule
+    @Rule
+    public static final DeviceState sDeviceState = new DeviceState();
+
+    private static final TestApis sTestApis = new TestApis();
+    private static final TestAppProvider sTestAppProvider = new TestAppProvider();
+    private static final TestApp sTestApp =
+            sTestAppProvider.query().whereActivities().isNotEmpty().get();
+    private RemoteDevicePolicyManager mDevicePolicyManager;
+    private DevicePolicyManager mLocalDevicePolicyManager;
+    private ComponentName mAdmin;
+    private UiAutomation mUiAutomation;
+    private static final int WAIT_IN_MILLISECOND = 500 * 20;
+
+
+    @Before
+    public void setUp() {
+        mAdmin = sDeviceState.dpc().componentName();
+        mDevicePolicyManager = sDeviceState.dpc().devicePolicyManager();
+        //TODO(b/198593716) : Use TestApi to take screnshot instead of UiAutomation.
+        mUiAutomation = InstrumentationRegistry.getInstrumentation().getUiAutomation();
+        mLocalDevicePolicyManager = ApplicationProvider.getApplicationContext().getSystemService(
+                DevicePolicyManager.class);
+    }
+
+    @After
+    public void tearDown() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, false);
+    }
+
+    @Test
+    @PositivePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_false_works() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, false);
+
+        assertThat(mLocalDevicePolicyManager.getScreenCaptureDisabled(/* admin= */ null)).isFalse();
+    }
+
+    @Test
+    @CanSetPolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_false_checkWithDPC_works() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, false);
+
+        assertThat(mDevicePolicyManager.getScreenCaptureDisabled(mAdmin)).isFalse();
+    }
+
+    @Test
+    @PositivePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_works() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+        assertThat(mLocalDevicePolicyManager.getScreenCaptureDisabled(/* admin= */ null)).isTrue();
+    }
+
+    @Test
+    @CanSetPolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_checkWithDPC_works() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+        assertThat(mDevicePolicyManager.getScreenCaptureDisabled(mAdmin)).isTrue();
+    }
+
+    @Test
+    @NegativePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_doesNotApply() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+        assertThat(mLocalDevicePolicyManager.getScreenCaptureDisabled(/* admin= */ null)).isFalse();
+    }
+
+    @Test
+    @NegativePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_screenCaptureWorks() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+        assertThat(takeScreenshotExpectingSuccess()).isNotNull();
+    }
+
+    @Test
+    @PositivePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_screenCaptureFails() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+        assertThat(takeScreenshotExpectingFailure()).isNull();
+    }
+
+    @Test
+    @PositivePolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_false_screenCaptureWorks() {
+        mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, false);
+
+        assertThat(takeScreenshotExpectingSuccess()).isNotNull();
+    }
+
+    @Test
+    @CanSetPolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_true_metricsLogged() {
+        try (EnterpriseMetricsRecorder metrics = EnterpriseMetricsRecorder.create()) {
+            mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, true);
+
+            assertThat(metrics.query()
+                    .whereType().isEqualTo(EventId.SET_SCREEN_CAPTURE_DISABLED_VALUE)
+                    .whereAdminPackageName().isEqualTo(mAdmin.getPackageName())
+                    .whereBoolean().isTrue()).wasLogged();
+        }
+    }
+
+    @Test
+    @CanSetPolicyTest(policy = ScreenCaptureDisabled.class)
+    @Postsubmit(reason = "new test")
+    public void setScreenCaptureDisabled_false_metricsLogged() {
+        try (EnterpriseMetricsRecorder metrics = EnterpriseMetricsRecorder.create()) {
+            mDevicePolicyManager.setScreenCaptureDisabled(mAdmin, false);
+
+            assertThat(metrics.query()
+                    .whereType().isEqualTo(EventId.SET_SCREEN_CAPTURE_DISABLED_VALUE)
+                    .whereAdminPackageName().isEqualTo(mAdmin.getPackageName())
+                    .whereBoolean().isFalse()).wasLogged();
+        }
+    }
+
+    private Bitmap takeScreenshotExpectingFailure() {
+        try (TestAppInstanceReference testApp = sTestApp.install(
+                sTestApis.users().instrumented())) {
+            testApp.activities().any().start();
+            return PollingCheck.waitFor(WAIT_IN_MILLISECOND, mUiAutomation::takeScreenshot,
+                    Objects::isNull);
+        }
+    }
+
+    private Bitmap takeScreenshotExpectingSuccess() {
+        try (TestAppInstanceReference testApp = sTestApp.install(
+                sTestApis.users().instrumented())) {
+            testApp.activities().any().start();
+            return PollingCheck.waitFor(WAIT_IN_MILLISECOND, mUiAutomation::takeScreenshot,
+                    Objects::nonNull);
+        }
+    }
+}