Use DesktopModeStatus#canEnterDesktopMode instead of just Flag check.

Test: Existing tests pass
Bug: 349544589
Flag: EXEMPT minor refactoring
Change-Id: I7050a16df1fe40684786ad5511415a2e07a457a6
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
index 271c07d..08de900 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
@@ -19,8 +19,6 @@
 import static android.view.WindowManager.LARGE_SCREEN_SMALLEST_SCREEN_WIDTH_DP;
 import static android.window.TaskConstants.TASK_CHILD_LAYER_COMPAT_UI;
 
-import static com.android.wm.shell.shared.desktopmode.DesktopModeFlags.DESKTOP_WINDOWING_MODE;
-
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.TaskInfo;
@@ -40,6 +38,7 @@
 import com.android.wm.shell.compatui.api.CompatUIEvent;
 import com.android.wm.shell.compatui.impl.CompatUIEvents.SizeCompatRestartButtonAppeared;
 import com.android.wm.shell.shared.desktopmode.DesktopModeFlags;
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 
 import java.util.function.Consumer;
 
@@ -83,7 +82,7 @@
         super(context, taskInfo, syncQueue, taskListener, displayLayout);
         mCallback = callback;
         mHasSizeCompat = taskInfo.appCompatTaskInfo.topActivityInSizeCompat;
-        if (DESKTOP_WINDOWING_MODE.isEnabled(mContext)
+        if (DesktopModeStatus.canEnterDesktopMode(context)
                 && DesktopModeFlags.DYNAMIC_INITIAL_BOUNDS.isEnabled(context)) {
             // Don't show the SCM button for freeform tasks
             mHasSizeCompat &= !taskInfo.isFreeform();
@@ -139,7 +138,7 @@
             boolean canShow) {
         final boolean prevHasSizeCompat = mHasSizeCompat;
         mHasSizeCompat = taskInfo.appCompatTaskInfo.topActivityInSizeCompat;
-        if (DESKTOP_WINDOWING_MODE.isEnabled(mContext)
+        if (DesktopModeStatus.canEnterDesktopMode(mContext)
                 && DesktopModeFlags.DYNAMIC_INITIAL_BOUNDS.isEnabled(mContext)) {
             // Don't show the SCM button for freeform tasks
             mHasSizeCompat &= !taskInfo.isFreeform();
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
index 3fe3162..29450a2 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
@@ -26,7 +26,7 @@
 import com.android.systemui.screenshot.policy.CapturePolicy.PolicyResult
 import com.android.systemui.screenshot.policy.CapturePolicy.PolicyResult.NotMatched
 import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
-import com.android.wm.shell.shared.desktopmode.DesktopModeFlags
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import javax.inject.Inject
 import kotlinx.coroutines.flow.first
 
@@ -48,7 +48,7 @@
             return NotMatched(policy = NAME, reason = SHADE_EXPANDED)
         }
 
-        if (DesktopModeFlags.DESKTOP_WINDOWING_MODE.isEnabled(context)) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
             content.rootTasks.firstOrNull()?.also {
                 if (it.windowingMode == WINDOWING_MODE_FREEFORM) {
                     return NotMatched(policy = NAME, reason = DESKTOP_MODE_ENABLED)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
index a5fbfb5..191a048 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
@@ -18,11 +18,13 @@
 
 import android.content.ComponentName
 import android.content.Context
+import android.content.res.Resources
 import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.internal.R
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.data.model.DisplayContentModel
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
@@ -47,6 +49,7 @@
 import com.android.systemui.screenshot.policy.WorkProfilePolicy.Companion.SHADE_EXPANDED
 import com.android.systemui.screenshot.policy.WorkProfilePolicy.Companion.WORK_TASK_IS_TOP
 import com.android.systemui.screenshot.policy.WorkProfilePolicy.Companion.WORK_TASK_NOT_TOP
+import com.android.systemui.util.mockito.whenever
 import com.android.window.flags.Flags
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
@@ -55,6 +58,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
+import org.mockito.MockitoAnnotations
 import org.mockito.junit.MockitoJUnit
 import org.mockito.junit.MockitoRule
 
@@ -66,12 +70,19 @@
     @JvmField @Rule(order = 2) val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
     @Mock lateinit var mContext: Context
+    @Mock lateinit var mResources: Resources
 
     private val kosmos = Kosmos()
     private lateinit var policy: WorkProfilePolicy
 
     @Before
     fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        // Set desktop mode supported
+        whenever(mContext.resources).thenReturn(mResources)
+        whenever(mResources.getBoolean(R.bool.config_isDesktopModeSupported)).thenReturn(true)
+
         policy = WorkProfilePolicy(kosmos.profileTypeRepository, mContext)
     }