Add shadows for getLockTaskFeatures and setLockTaskFeatures.

PiperOrigin-RevId: 487133948
diff --git a/robolectric/src/test/java/org/robolectric/shadows/ShadowDevicePolicyManagerTest.java b/robolectric/src/test/java/org/robolectric/shadows/ShadowDevicePolicyManagerTest.java
index 09a3b54..7d36f35 100644
--- a/robolectric/src/test/java/org/robolectric/shadows/ShadowDevicePolicyManagerTest.java
+++ b/robolectric/src/test/java/org/robolectric/shadows/ShadowDevicePolicyManagerTest.java
@@ -9,6 +9,9 @@
 import static android.app.admin.DevicePolicyManager.ENCRYPTION_STATUS_ACTIVE_PER_USER;
 import static android.app.admin.DevicePolicyManager.ENCRYPTION_STATUS_INACTIVE;
 import static android.app.admin.DevicePolicyManager.ENCRYPTION_STATUS_UNSUPPORTED;
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_HOME;
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS;
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_OVERVIEW;
 import static android.app.admin.DevicePolicyManager.PASSWORD_COMPLEXITY_HIGH;
 import static android.app.admin.DevicePolicyManager.PERMISSION_POLICY_AUTO_GRANT;
 import static android.app.admin.DevicePolicyManager.STATE_USER_SETUP_COMPLETE;
@@ -1698,6 +1701,91 @@
   }
 
   @Test
+  @Config(minSdk = P)
+  public void getLockTaskFeatures_nullAdmin_throwsNullPointerException() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+    assertThrows(NullPointerException.class, () -> devicePolicyManager.getLockTaskFeatures(null));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void getLockTaskFeatures_notOwner_throwsSecurityException() {
+    assertThrows(
+        SecurityException.class, () -> devicePolicyManager.getLockTaskFeatures(testComponent));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void getLockTaskFeatures_default_noFeatures() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+
+    assertThat(devicePolicyManager.getLockTaskFeatures(testComponent)).isEqualTo(0);
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_nullAdmin_throwsNullPointerException() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+
+    assertThrows(
+        NullPointerException.class, () -> devicePolicyManager.setLockTaskFeatures(null, 0));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_notOwner_throwsSecurityException() {
+    assertThrows(
+        SecurityException.class, () -> devicePolicyManager.setLockTaskFeatures(testComponent, 0));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_overviewWithoutHome_throwsIllegalArgumentException() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> devicePolicyManager.setLockTaskFeatures(testComponent, LOCK_TASK_FEATURE_OVERVIEW));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_notificationsWithoutHome_throwsIllegalArgumentException() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
+            devicePolicyManager.setLockTaskFeatures(
+                testComponent, LOCK_TASK_FEATURE_NOTIFICATIONS));
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_homeOverviewNotifications_success() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+
+    int flags =
+        LOCK_TASK_FEATURE_HOME | LOCK_TASK_FEATURE_OVERVIEW | LOCK_TASK_FEATURE_NOTIFICATIONS;
+    devicePolicyManager.setLockTaskFeatures(testComponent, flags);
+
+    assertThat(devicePolicyManager.getLockTaskFeatures(testComponent)).isEqualTo(flags);
+  }
+
+  @Test
+  @Config(minSdk = P)
+  public void setLockTaskFeatures_setFeaturesTwice_keepsLatestFeatures() {
+    shadowOf(devicePolicyManager).setProfileOwner(testComponent);
+    devicePolicyManager.setLockTaskFeatures(testComponent, LOCK_TASK_FEATURE_HOME);
+
+    int flags =
+        LOCK_TASK_FEATURE_HOME | LOCK_TASK_FEATURE_OVERVIEW | LOCK_TASK_FEATURE_NOTIFICATIONS;
+    devicePolicyManager.setLockTaskFeatures(testComponent, flags);
+
+    assertThat(devicePolicyManager.getLockTaskFeatures(testComponent)).isEqualTo(flags);
+  }
+
+  @Test
   @Config(minSdk = LOLLIPOP)
   public void getLockTaskPackages_notOwner() {
     try {
diff --git a/shadows/framework/src/main/java/org/robolectric/shadows/ShadowDevicePolicyManager.java b/shadows/framework/src/main/java/org/robolectric/shadows/ShadowDevicePolicyManager.java
index c768882..dbaa353 100644
--- a/shadows/framework/src/main/java/org/robolectric/shadows/ShadowDevicePolicyManager.java
+++ b/shadows/framework/src/main/java/org/robolectric/shadows/ShadowDevicePolicyManager.java
@@ -1,5 +1,8 @@
 package org.robolectric.shadows;
 
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_HOME;
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS;
+import static android.app.admin.DevicePolicyManager.LOCK_TASK_FEATURE_OVERVIEW;
 import static android.os.Build.VERSION_CODES.JELLY_BEAN_MR1;
 import static android.os.Build.VERSION_CODES.JELLY_BEAN_MR2;
 import static android.os.Build.VERSION_CODES.LOLLIPOP;
@@ -128,6 +131,7 @@
   private final Map<ComponentName, CharSequence> longSupportMessageMap = new HashMap<>();
   private final Set<ComponentName> componentsWithActivatedTokens = new HashSet<>();
   private Collection<String> packagesToFailForSetApplicationHidden = Collections.emptySet();
+  private int lockTaskFeatures;
   private final List<String> lockTaskPackages = new ArrayList<>();
   private Context context;
   private ApplicationPackageManager applicationPackageManager;
@@ -1236,6 +1240,31 @@
     return policyGrantedSet != null && policyGrantedSet.contains(usesPolicy);
   }
 
+  @Implementation(minSdk = P)
+  protected int getLockTaskFeatures(ComponentName admin) {
+    Objects.requireNonNull(admin, "ComponentName is null");
+    enforceDeviceOwnerOrProfileOwner(admin);
+    return lockTaskFeatures;
+  }
+
+  @Implementation(minSdk = P)
+  protected void setLockTaskFeatures(ComponentName admin, int flags) {
+    Objects.requireNonNull(admin, "ComponentName is null");
+    enforceDeviceOwnerOrProfileOwner(admin);
+    // Throw if Overview is used without Home.
+    boolean hasHome = (flags & LOCK_TASK_FEATURE_HOME) != 0;
+    boolean hasOverview = (flags & LOCK_TASK_FEATURE_OVERVIEW) != 0;
+    Preconditions.checkArgument(
+        hasHome || !hasOverview,
+        "Cannot use LOCK_TASK_FEATURE_OVERVIEW without LOCK_TASK_FEATURE_HOME");
+    boolean hasNotification = (flags & LOCK_TASK_FEATURE_NOTIFICATIONS) != 0;
+    Preconditions.checkArgument(
+        hasHome || !hasNotification,
+        "Cannot use LOCK_TASK_FEATURE_NOTIFICATIONS without LOCK_TASK_FEATURE_HOME");
+
+    lockTaskFeatures = flags;
+  }
+
   @Implementation(minSdk = LOLLIPOP)
   protected void setLockTaskPackages(@NonNull ComponentName admin, String[] packages) {
     enforceDeviceOwnerOrProfileOwner(admin);