diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleScreen.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleScreen.kt
index 0a3321eb5a721..a57d3e63b559f 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleScreen.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/bubble/LaunchBubbleScreen.kt
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.flicker.bubble
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
index 19e020a52deec..5fc80db6c6176 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/ExpandPipOnDoubleClickTest.kt
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.flicker.pip
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.view.Surface
 import androidx.test.filters.RequiresDevice
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest.kt
index 8729bb6776f0b..87e927fd50ea6 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest.kt
@@ -24,7 +24,10 @@ import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
 import com.android.server.wm.flicker.annotation.Group3
 import com.android.server.wm.flicker.dsl.FlickerBuilder
+import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import com.android.server.wm.flicker.traces.region.RegionSubject
+import org.junit.Assume
+import org.junit.Before
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -58,6 +61,11 @@ import org.junit.runners.Parameterized
 open class MovePipDownShelfHeightChangeTest(
     testSpec: FlickerTestParameter
 ) : MovePipShelfHeightTransition(testSpec) {
+    @Before
+    open fun before() {
+        Assume.assumeFalse(isShellTransitionsEnabled)
+    }
+
     /**
      * Defines the transition used to run the test
      */
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest_ShellTransit.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest_ShellTransit.kt
index 6c80daadfdd94..0ff260b94dc84 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest_ShellTransit.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipDownShelfHeightChangeTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.flicker.pip
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -58,7 +58,7 @@ class MovePipDownShelfHeightChangeTest_ShellTransit(
     testSpec: FlickerTestParameter
 ) : MovePipDownShelfHeightChangeTest(testSpec) {
     @Before
-    fun before() {
+    override fun before() {
         Assume.assumeTrue(isShellTransitionsEnabled)
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipUpShelfHeightChangeTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipUpShelfHeightChangeTest.kt
index 4bc8eb13e26f0..388b5e0b5e47c 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipUpShelfHeightChangeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/MovePipUpShelfHeightChangeTest.kt
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.flicker.pip
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import android.view.Surface
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
index 3b3a3035f2960..f8348203599fe 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest_ShellTransit.kt
index 849e3ae91eaa8..0d2869cbd178c 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppColdTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
index 18d017d95d035..00fee82adf764 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest.kt
@@ -16,10 +16,10 @@
 
 package com.android.server.wm.flicker.launch
 
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
 import android.view.Display
-import androidx.test.filters.FlakyTest
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
 import com.android.server.wm.flicker.FlickerTestParameterFactory
@@ -31,7 +31,6 @@ import com.android.server.wm.flicker.helpers.reopenAppFromOverview
 import com.android.server.wm.flicker.helpers.setRotation
 import com.android.server.wm.traces.common.WindowManagerConditionsFactory
 import org.junit.Assume.assumeFalse
-import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.FixMethodOrder
 import org.junit.Test
@@ -140,37 +139,6 @@ open class OpenAppFromOverviewTest(testSpec: FlickerTestParameter)
     @Test
     override fun appWindowBecomesVisible() = super.appWindowBecomesVisible_warmStart()
 
-    /** {@inheritDoc} */
-    @Presubmit
-    @Test
-    override fun appWindowReplacesLauncherAsTopWindow() {
-        assumeFalse(isShellTransitionsEnabled)
-        super.appWindowReplacesLauncherAsTopWindow()
-    }
-
-    @FlakyTest(bugId = 216266712)
-    @Test
-    fun appWindowReplacesLauncherAsTopWindow_shellTransit() {
-        assumeTrue(isShellTransitionsEnabled)
-        super.appWindowReplacesLauncherAsTopWindow()
-    }
-
-    /** {@inheritDoc} */
-    @Presubmit
-    @Test
-    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() {
-        assumeFalse(isShellTransitionsEnabled)
-        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
-    }
-
-    /** {@inheritDoc} */
-    @FlakyTest(bugId = 218470989)
-    @Test
-    fun visibleWindowsShownMoreThanOneConsecutiveEntry_shellTransit() {
-        assumeTrue(isShellTransitionsEnabled)
-        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
-    }
-
     companion object {
         /**
          * Creates the test configurations.
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest_ShellTransit.kt
index 24716ffb902a6..1c06495cf4730 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppFromOverviewTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
@@ -25,6 +25,7 @@ import com.android.server.wm.flicker.helpers.isShellTransitionsEnabled
 import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.FixMethodOrder
+import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.MethodSorters
 import org.junit.runners.Parameterized
@@ -60,4 +61,16 @@ class OpenAppFromOverviewTest_ShellTransit(testSpec: FlickerTestParameter)
     override fun before() {
         assumeTrue(isShellTransitionsEnabled)
     }
+
+    /** {@inheritDoc} */
+    @FlakyTest(bugId = 216266712)
+    @Test
+    override fun appWindowReplacesLauncherAsTopWindow() =
+            super.appWindowReplacesLauncherAsTopWindow()
+
+    /** {@inheritDoc} */
+    @FlakyTest(bugId = 218470989)
+    @Test
+    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
+        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
 }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
index 0ba5369917b8c..b0e53e9f85bf1 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.platform.test.annotations.RequiresDevice
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest_ShellTransit.kt
index f8afcd9233f40..8a08d07e654e0 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppNonResizeableTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
index d2f6c7f144dd4..53560cc14c44b 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppTransition.kt
@@ -17,7 +17,7 @@
 package com.android.server.wm.flicker.launch
 
 import android.app.Instrumentation
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.Presubmit
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.server.wm.flicker.FlickerBuilderProvider
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest_ShellTransit.kt
index 0ce73f4e25092..3958dd2c89b4a 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/launch/OpenAppWarmTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.launch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest_ShellTransit.kt
index 07fe274e5ea20..cffed81f6cb96 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/quickswitch/QuickSwitchBetweenTwoAppsBackTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.quickswitch
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/ChangeAppRotationTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/ChangeAppRotationTest.kt
index f603f6e7ed9d8..2b944c6667823 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/ChangeAppRotationTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/ChangeAppRotationTest.kt
@@ -112,10 +112,7 @@ class ChangeAppRotationTest(
      * Checks that the [FlickerComponentName.ROTATION] layer appears during the transition,
      * doesn't flicker, and disappears before the transition is complete
      */
-    @Presubmit
-    @Test
-    fun rotationLayerAppearsAndVanishes() {
-        Assume.assumeFalse(isShellTransitionsEnabled)
+    fun rotationLayerAppearsAndVanishesAssertion() {
         testSpec.assertLayers {
             this.isVisible(testApp.component)
                 .then()
@@ -126,11 +123,26 @@ class ChangeAppRotationTest(
         }
     }
 
+    /**
+     * Checks that the [FlickerComponentName.ROTATION] layer appears during the transition,
+     * doesn't flicker, and disappears before the transition is complete
+     */
+    @Presubmit
+    @Test
+    fun rotationLayerAppearsAndVanishes() {
+        Assume.assumeFalse(isShellTransitionsEnabled)
+        rotationLayerAppearsAndVanishesAssertion()
+    }
+
+    /**
+     * Checks that the [FlickerComponentName.ROTATION] layer appears during the transition,
+     * doesn't flicker, and disappears before the transition is complete
+     */
     @FlakyTest(bugId = 218484127)
     @Test
     fun rotationLayerAppearsAndVanishes_shellTransit() {
         Assume.assumeTrue(isShellTransitionsEnabled)
-        rotationLayerAppearsAndVanishes()
+        rotationLayerAppearsAndVanishesAssertion()
     }
 
     /**
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
index d1bdeed81b788..0becadf630e16 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/RotationTransition.kt
@@ -153,4 +153,4 @@ abstract class RotationTransition(protected val testSpec: FlickerTestParameter)
             }
         }
     }
-}
\ No newline at end of file
+}
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/SeamlessAppRotationTest_ShellTransit.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/SeamlessAppRotationTest_ShellTransit.kt
index be55751bf0600..d397d5979803a 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/SeamlessAppRotationTest_ShellTransit.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/rotation/SeamlessAppRotationTest_ShellTransit.kt
@@ -16,7 +16,7 @@
 
 package com.android.server.wm.flicker.rotation
 
-import android.platform.test.annotations.FlakyTest
+import androidx.test.filters.FlakyTest
 import android.platform.test.annotations.RequiresDevice
 import com.android.server.wm.flicker.FlickerParametersRunnerFactory
 import com.android.server.wm.flicker.FlickerTestParameter