diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
index 2ecf01f6999d0c881667deafa38ea3d14ae371ad..1237347428206d11603a38ffabc6c675111a2f72 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
@@ -16,7 +16,8 @@
 
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
-import android.platform.test.flag.junit.SetFlagsRule
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_OTHER_DEVICE_CONNECTION
@@ -60,8 +61,6 @@ class WifiViewModelTest : SysuiTestCase() {
 
     private lateinit var underTest: WifiViewModel
 
-    private val setFlagsRule = SetFlagsRule()
-
     @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
@@ -187,11 +186,9 @@ class WifiViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @DisableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun activity_nullSsid_outputsFalse_staticFlagOff() =
         testScope.runTest {
-            // GIVEN flag is disabled
-            setFlagsRule.disableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             whenever(connectivityConstants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()
 
@@ -214,11 +211,9 @@ class WifiViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @EnableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun activity_nullSsid_outputsFalse_staticFlagOn() =
         testScope.runTest {
-            // GIVEN flag is enabled
-            setFlagsRule.enableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             whenever(connectivityConstants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()
 
@@ -371,11 +366,9 @@ class WifiViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @DisableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun activityContainer_inAndOutFalse_outputsTrue_staticFlagOff() =
         testScope.runTest {
-            // GIVEN the flag is off
-            setFlagsRule.disableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             whenever(connectivityConstants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()
             wifiRepository.setWifiNetwork(ACTIVE_VALID_WIFI_NETWORK)
@@ -389,11 +382,9 @@ class WifiViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @EnableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun activityContainer_inAndOutFalse_outputsTrue_staticFlagOn() =
         testScope.runTest {
-            // GIVEN the flag is on
-            setFlagsRule.enableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             whenever(connectivityConstants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()
             wifiRepository.setWifiNetwork(ACTIVE_VALID_WIFI_NETWORK)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
index 0f779d977c71ce1351f8a050cd46f27b7229b37b..44fa132839912b4e6bdd141a971176211bbfd640 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
@@ -16,7 +16,8 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.ui.viewmodel
 
-import android.platform.test.flag.junit.SetFlagsRule
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.AccessibilityContentDescriptions.PHONE_SIGNAL_STRENGTH
@@ -73,8 +74,6 @@ import org.mockito.MockitoAnnotations
 class MobileIconViewModelTest : SysuiTestCase() {
     private var connectivityRepository = FakeConnectivityRepository()
 
-    private val setFlagsRule = SetFlagsRule()
-
     private lateinit var underTest: MobileIconViewModel
     private lateinit var interactor: MobileIconInteractorImpl
     private lateinit var iconsInteractor: MobileIconsInteractorImpl
@@ -561,11 +560,9 @@ class MobileIconViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @DisableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun dataActivity_configOn_testIndicators_staticFlagOff() =
         testScope.runTest {
-            // GIVEN STATUS_BAR_STATIC_NETWORK_INDICATORS flag is off
-            setFlagsRule.disableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             // Create a new view model here so the constants are properly read
             whenever(constants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()
@@ -618,11 +615,9 @@ class MobileIconViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    @EnableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
     fun dataActivity_configOn_testIndicators_staticFlagOn() =
         testScope.runTest {
-            // GIVEN STATUS_BAR_STATIC_NETWORK_INDICATORS flag is on
-            setFlagsRule.enableFlags(FLAG_STATUS_BAR_STATIC_INOUT_INDICATORS)
-
             // Create a new view model here so the constants are properly read
             whenever(constants.shouldShowActivityConfig).thenReturn(true)
             createAndSetViewModel()