Create CompatUIShellTestCase

Create a base class for CompatUI tests in Shell with common
rules about Flags.

Flag: EXEMPT Refactoring
Fix: 351114642
Test: atest WMShellUnitTests

Change-Id: I216d709d7e21e91ca625c4ba4d7571a68c27a98e

14 files changed, 55 insertions, 119 deletions

diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/AppCompatUtilsTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/AppCompatUtilsTest.kt
index d52fd4fdf6c7..3ab7b3418e02 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/AppCompatUtilsTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/AppCompatUtilsTest.kt
@@ -20,7 +20,6 @@ import android.content.ComponentName
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.internal.R
-import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
@@ -35,7 +34,7 @@ import org.junit.runner.RunWith
  */
 @RunWith(AndroidTestingRunner::class)
 @SmallTest
-class AppCompatUtilsTest : ShellTestCase() {
+class AppCompatUtilsTest : CompatUIShellTestCase() {
     @Test
     fun testIsTopActivityExemptFromDesktopWindowing_onlyTransparentActivitiesInStack() {
         assertTrue(isTopActivityExemptFromDesktopWindowing(mContext,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
index 67573dabf2ec..ecf766db0be3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
@@ -39,9 +39,6 @@ import android.content.res.Configuration;
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.view.InsetsSource;
 import android.view.InsetsState;
@@ -52,7 +49,6 @@ import androidx.test.filters.SmallTest;
 
 import com.android.window.flags.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
 import com.android.wm.shell.common.DisplayInsetsController;
@@ -70,11 +66,8 @@ import com.android.wm.shell.transition.Transitions;
 
 import dagger.Lazy;
 
-import java.util.Optional;
-
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -82,25 +75,20 @@ import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.Optional;
+
 /**
  * Tests for {@link CompatUIController}.
  *
  * Build/Install/Run:
- *  atest WMShellUnitTests:CompatUIControllerTest
+ * atest WMShellUnitTests:CompatUIControllerTest
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class CompatUIControllerTest extends ShellTestCase {
+public class CompatUIControllerTest extends CompatUIShellTestCase {
     private static final int DISPLAY_ID = 0;
     private static final int TASK_ID = 12;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
-
     private CompatUIController mController;
     private ShellInit mShellInit;
     @Mock
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
index e5d1919decf4..2117b062bf57 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
@@ -26,8 +26,6 @@ import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.graphics.Rect;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.LayoutInflater;
@@ -40,7 +38,6 @@ import androidx.test.filters.SmallTest;
 import com.android.window.flags.Flags;
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
@@ -49,7 +46,6 @@ import com.android.wm.shell.compatui.api.CompatUIEvent;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -66,14 +62,10 @@ import java.util.function.Consumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class CompatUILayoutTest extends ShellTestCase {
+public class CompatUILayoutTest extends CompatUIShellTestCase {
 
     private static final int TASK_ID = 1;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
     @Mock private Consumer<CompatUIEvent> mCallback;
     @Mock private Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> mOnRestartButtonClicked;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIShellTestCase.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIShellTestCase.java
new file mode 100644
index 000000000000..5a49d01f2991
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIShellTestCase.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.compatui;
+
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import com.android.wm.shell.ShellTestCase;
+
+import org.junit.Rule;
+
+/**
+ * Base class for CompatUI tests.
+ */
+public class CompatUIShellTestCase extends ShellTestCase {
+
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIStatusManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIStatusManagerTest.java
index 8fd7c0ec3099..0b37648faeec 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIStatusManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIStatusManagerTest.java
@@ -27,8 +27,6 @@ import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.wm.shell.ShellTestCase;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -44,7 +42,7 @@ import java.util.function.IntSupplier;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class CompatUIStatusManagerTest extends ShellTestCase {
+public class CompatUIStatusManagerTest extends CompatUIShellTestCase {
 
     private FakeCompatUIStatusManagerTest mTestState;
     private CompatUIStatusManager mStatusManager;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
index 1c0175603df9..61b6d803c8be 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
@@ -16,7 +16,6 @@
 
 package com.android.wm.shell.compatui;
 
-import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
 import static android.view.WindowInsets.Type.navigationBars;
 import static android.view.WindowManager.LARGE_SCREEN_SMALLEST_SCREEN_WIDTH_DP;
 
@@ -40,9 +39,6 @@ import android.app.TaskInfo;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.DisplayInfo;
@@ -56,7 +52,6 @@ import androidx.test.filters.SmallTest;
 
 import com.android.window.flags.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
@@ -65,7 +60,6 @@ import com.android.wm.shell.compatui.api.CompatUIEvent;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -82,13 +76,7 @@ import java.util.function.Consumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class CompatUIWindowManagerTest extends ShellTestCase {
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
-
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
+public class CompatUIWindowManagerTest extends CompatUIShellTestCase {
 
     private static final int TASK_ID = 1;
     private static final int TASK_WIDTH = 2000;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
index e8191db13084..e786fef1855c 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
@@ -25,8 +25,6 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -34,10 +32,8 @@ import android.view.View;
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.R;
-import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -51,7 +47,7 @@ import org.mockito.MockitoAnnotations;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class LetterboxEduDialogLayoutTest extends ShellTestCase {
+public class LetterboxEduDialogLayoutTest extends CompatUIShellTestCase {
 
     @Mock
     private Runnable mDismissCallback;
@@ -60,10 +56,6 @@ public class LetterboxEduDialogLayoutTest extends ShellTestCase {
     private View mDismissButton;
     private View mDialogContainer;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
index 4c97c76ae122..09fc082a63e3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
@@ -46,9 +46,6 @@ import android.graphics.Rect;
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.DisplayCutout;
@@ -65,7 +62,6 @@ import androidx.test.filters.SmallTest;
 import com.android.window.flags.Flags;
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.DockStateReader;
@@ -75,7 +71,6 @@ import com.android.wm.shell.transition.Transitions;
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -95,7 +90,7 @@ import java.util.function.Consumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class LetterboxEduWindowManagerTest extends ShellTestCase {
+public class LetterboxEduWindowManagerTest extends CompatUIShellTestCase {
 
     private static final int USER_ID_1 = 1;
     private static final int USER_ID_2 = 2;
@@ -128,18 +123,11 @@ public class LetterboxEduWindowManagerTest extends ShellTestCase {
     @Mock private Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> mOnDismissCallback;
     @Mock private DockStateReader mDockStateReader;
 
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
-
     private CompatUIConfiguration mCompatUIConfiguration;
     private TestShellExecutor mExecutor;
     private FakeCompatUIStatusManagerTest mCompatUIStatus;
     private CompatUIStatusManager mCompatUIStatusManager;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
index 0da14d673732..02c099b3cfb2 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
@@ -26,8 +26,6 @@ import static org.mockito.Mockito.verify;
 
 import android.app.TaskInfo;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.LayoutInflater;
@@ -36,10 +34,8 @@ import android.view.View;
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.R;
-import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -54,7 +50,7 @@ import org.mockito.MockitoAnnotations;
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
-public class ReachabilityEduLayoutTest extends ShellTestCase {
+public class ReachabilityEduLayoutTest extends CompatUIShellTestCase {
 
     private ReachabilityEduLayout mLayout;
     private View mMoveUpButton;
@@ -68,10 +64,6 @@ public class ReachabilityEduLayoutTest extends ShellTestCase {
     @Mock
     private TaskInfo mTaskInfo;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
index eafb41470cda..fa04e070250e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
@@ -25,14 +25,11 @@ import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.res.Configuration;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
@@ -40,7 +37,6 @@ import com.android.wm.shell.common.SyncTransactionQueue;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -56,7 +52,7 @@ import java.util.function.BiConsumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class ReachabilityEduWindowManagerTest extends ShellTestCase {
+public class ReachabilityEduWindowManagerTest extends CompatUIShellTestCase {
     @Mock
     private SyncTransactionQueue mSyncTransactionQueue;
     @Mock
@@ -71,10 +67,6 @@ public class ReachabilityEduWindowManagerTest extends ShellTestCase {
     private TaskInfo mTaskInfo;
     private ReachabilityEduWindowManager mWindowManager;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
index 6b0c5dd2e1c7..2cded9d9776c 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
@@ -26,8 +26,6 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -36,10 +34,8 @@ import android.widget.CheckBox;
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.R;
-import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -55,7 +51,7 @@ import java.util.function.Consumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class RestartDialogLayoutTest extends ShellTestCase  {
+public class RestartDialogLayoutTest extends CompatUIShellTestCase  {
 
     @Mock private Runnable mDismissCallback;
     @Mock private Consumer<Boolean> mRestartCallback;
@@ -66,10 +62,6 @@ public class RestartDialogLayoutTest extends ShellTestCase  {
     private View mDialogContainer;
     private CheckBox mDontRepeatCheckBox;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
index cfeef90cb4b6..ebd0f412a0a1 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
@@ -22,15 +22,12 @@ import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.res.Configuration;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.transition.Transitions;
@@ -38,7 +35,6 @@ import com.android.wm.shell.transition.Transitions;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -54,7 +50,7 @@ import java.util.function.Consumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class RestartDialogWindowManagerTest extends ShellTestCase {
+public class RestartDialogWindowManagerTest extends CompatUIShellTestCase {
 
     @Mock
     private SyncTransactionQueue mSyncTransactionQueue;
@@ -66,10 +62,6 @@ public class RestartDialogWindowManagerTest extends ShellTestCase {
     private RestartDialogWindowManager mWindowManager;
     private TaskInfo mTaskInfo;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
index e8e68bdbd940..c6532e13f3cc 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
@@ -27,8 +27,6 @@ import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.ComponentName;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.LayoutInflater;
@@ -40,7 +38,6 @@ import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
@@ -48,7 +45,6 @@ import com.android.wm.shell.common.SyncTransactionQueue;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -66,7 +62,7 @@ import java.util.function.BiConsumer;
  */
 @RunWith(AndroidTestingRunner.class)
 @SmallTest
-public class UserAspectRatioSettingsLayoutTest extends ShellTestCase {
+public class UserAspectRatioSettingsLayoutTest extends CompatUIShellTestCase {
 
     private static final int TASK_ID = 1;
 
@@ -88,10 +84,6 @@ public class UserAspectRatioSettingsLayoutTest extends ShellTestCase {
     private UserAspectRatioSettingsLayout mLayout;
     private TaskInfo mTaskInfo;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
index 9f86d49b52c4..096e900199ba 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
@@ -41,8 +41,6 @@ import android.content.Intent;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper.RunWithLooper;
 import android.util.Pair;
@@ -56,7 +54,6 @@ import android.view.View;
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
@@ -65,7 +62,6 @@ import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
 import junit.framework.Assert;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -87,7 +83,7 @@ import java.util.function.Supplier;
 @RunWith(AndroidTestingRunner.class)
 @RunWithLooper
 @SmallTest
-public class UserAspectRatioSettingsWindowManagerTest extends ShellTestCase {
+public class UserAspectRatioSettingsWindowManagerTest extends CompatUIShellTestCase {
 
     private static final int TASK_ID = 1;
 
@@ -112,10 +108,6 @@ public class UserAspectRatioSettingsWindowManagerTest extends ShellTestCase {
 
     private TestShellExecutor mExecutor;
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule =
-            DeviceFlagsValueProvider.createCheckFlagsRule();
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);