diff --git a/res/values/strings.xml b/res/values/strings.xml
index 0c04ea8..c71b4db 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -11982,7 +11982,7 @@
     <!-- Toast on failure to reformat data to ext4 -->
     <string name="format_ext4_failure_toast">Failed to reformat and wipe the data partition to ext4.</string>
     <!-- Dialog to OEM unlock the device before using 16K developer option -->
-    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required for 16KB Mode</string>
+    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required</string>
     <string name="confirm_oem_unlock_for_16k_text">This device needs to have the bootloader unlocked before using the 16KB developer option.
         Software integrity cannot be guaranteed in this mode, and any data stored on the phone while the bootloader is unlocked may be at risk.
         All user data and settings will be wiped when activating 16KB mode. Once the bootloader is unlocked, activating the 16KB option will require two reboots.
@@ -12011,10 +12011,10 @@
     <!-- persistent notification 16k page agnostic mode text -->
     <string name="page_agnostic_16k_pages_text_short">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. Tap to read more.</string>
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. Tap to read more.</string>
     <string name="page_agnostic_16k_pages_text">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. This would factory reset the device again and
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. This would factory reset the device again and
         restore it to production settings. After the device successfully boots into Android, disable OEM unlocking in Developer options.
         If the device fails to boot into Android or is unstable, re-flash the device with the latest factory images from
         &lt;a href=\"https://developers.google.com/android/images\"&gt;https://developers.google.com/android/images&lt;/a&gt;
diff --git a/src/com/android/settings/development/Enable16KBootReceiver.java b/src/com/android/settings/development/Enable16KBootReceiver.java
index 007a67b..bc2096d 100644
--- a/src/com/android/settings/development/Enable16KBootReceiver.java
+++ b/src/com/android/settings/development/Enable16KBootReceiver.java
@@ -28,17 +28,17 @@
     @Override
     public void onReceive(@NonNull Context context, @NonNull Intent intent) {
         String action = intent.getAction();
-        if (!Intent.ACTION_BOOT_COMPLETED.equals(action)) {
-            return;
-        }
+        if (Intent.ACTION_BOOT_COMPLETED.equals(action)
+                || PageAgnosticNotificationService.INTENT_ACTION_DISMISSED.equals(action)) {
+            // Do nothing if device is not in page-agnostic mode
+            if (!Enable16kUtils.isPageAgnosticModeOn(context)) {
+                return;
+            }
 
-        // Do nothing if device is not in page-agnostic mode
-        if (!Enable16kUtils.isPageAgnosticModeOn(context)) {
-            return;
+            // start a service to post persistent notification
+            Intent startServiceIntent = new Intent(context, PageAgnosticNotificationService.class);
+            startServiceIntent.setAction(action);
+            context.startServiceAsUser(startServiceIntent, UserHandle.SYSTEM);
         }
-
-        // start a service to post persistent notification
-        Intent startNotificationIntent = new Intent(context, PageAgnosticNotificationService.class);
-        context.startServiceAsUser(startNotificationIntent, UserHandle.SYSTEM);
     }
 }
diff --git a/src/com/android/settings/development/OWNERS b/src/com/android/settings/development/OWNERS
index 92c8cfa..cc573ea 100644
--- a/src/com/android/settings/development/OWNERS
+++ b/src/com/android/settings/development/OWNERS
@@ -9,3 +9,6 @@
 
 # ShowRefreshRatePreferenceController
 per-file ShowRefreshRatePreferenceController.java=file:platform/frameworks/native:/services/surfaceflinger/OWNERS
+
+# DesktopModePreferenceController
+per-file DesktopModePreferenceController.java=file:platform/frameworks/base:/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/OWNERS
diff --git a/src/com/android/settings/development/PageAgnosticNotificationService.java b/src/com/android/settings/development/PageAgnosticNotificationService.java
index bce1dd9..d33f1da 100644
--- a/src/com/android/settings/development/PageAgnosticNotificationService.java
+++ b/src/com/android/settings/development/PageAgnosticNotificationService.java
@@ -34,6 +34,8 @@
 
     private static final String NOTIFICATION_CHANNEL_ID =
             "com.android.settings.development.PageAgnosticNotificationService";
+    public static final String INTENT_ACTION_DISMISSED =
+            "com.android.settings.development.NOTIFICATION_DISMISSED";
     private static final int NOTIFICATION_ID = 1;
 
     static final int DISABLE_UPDATES_SETTING = 1;
@@ -63,6 +65,9 @@
     public void onCreate() {
         super.onCreate();
         createNotificationChannel();
+
+        // No updates should be allowed in page-agnostic mode
+        disableAutomaticUpdates();
     }
 
     private Notification buildNotification() {
@@ -89,6 +94,15 @@
                         notifyIntent,
                         PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
 
+        Intent dismissIntent = new Intent(this, Enable16KBootReceiver.class);
+        dismissIntent.setAction(INTENT_ACTION_DISMISSED);
+        PendingIntent dismissPendingIntent =
+                PendingIntent.getBroadcast(
+                        this.getApplicationContext(),
+                        0,
+                        dismissIntent,
+                        PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
+
         Notification.Action action =
                 new Notification.Action.Builder(
                                 R.drawable.empty_icon,
@@ -96,14 +110,15 @@
                                 notifyPendingIntent)
                         .build();
 
+        // TODO:(b/349860833) Change text style to BigTextStyle once the ellipsis issue is fixed.
         Notification.Builder builder =
                 new Notification.Builder(this, NOTIFICATION_CHANNEL_ID)
                         .setContentTitle(title)
                         .setContentText(text)
                         .setOngoing(true)
                         .setSmallIcon(R.drawable.ic_settings_24dp)
-                        .setStyle(new Notification.BigTextStyle().bigText(text))
                         .setContentIntent(notifyPendingIntent)
+                        .setDeleteIntent(dismissPendingIntent)
                         .addAction(action);
 
         return builder.build();
@@ -131,9 +146,6 @@
         if (mNotificationManager != null) {
             mNotificationManager.notify(NOTIFICATION_ID, notification);
         }
-
-        // No updates should be allowed in page-agnostic mode
-        disableAutomaticUpdates();
-        return Service.START_NOT_STICKY;
+        return Service.START_REDELIVER_INTENT;
     }
 }
diff --git a/src/com/android/settings/sim/smartForwarding/DisableSmartForwardingTask.java b/src/com/android/settings/sim/smartForwarding/DisableSmartForwardingTask.java
index a1035dc..8448319 100644
--- a/src/com/android/settings/sim/smartForwarding/DisableSmartForwardingTask.java
+++ b/src/com/android/settings/sim/smartForwarding/DisableSmartForwardingTask.java
@@ -23,6 +23,14 @@
 import android.telephony.TelephonyManager;
 import android.util.Log;
 
+import com.android.settings.sim.smartForwarding.EnableSmartForwardingTask.UpdateCommand;
+
+import com.google.common.util.concurrent.SettableFuture;
+
+import java.util.ArrayList;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+
 public class DisableSmartForwardingTask implements Runnable {
     private final TelephonyManager tm;
     private final boolean[] callWaitingStatus;
@@ -37,22 +45,123 @@
 
     @Override
     public void run() {
-        for (int i = 0; i < tm.getActiveModemCount(); i++) {
-            int subId = getSubId(i);
-            if (callWaitingStatus != null
-                    && subId != SubscriptionManager.INVALID_SUBSCRIPTION_ID) {
-                Log.d(TAG, "Restore call waiting to " + callWaitingStatus[i]);
-                tm.createForSubscriptionId(subId)
-                        .setCallWaitingEnabled(callWaitingStatus[i], null, null);
+        FlowController controller = new FlowController();
+        if (controller.init()) {
+            controller.startProcess();
+        }
+    }
+
+    class FlowController {
+        private final ArrayList<UpdateCommand> mSteps = new ArrayList<>();
+
+        /* package */ boolean init() {
+            if (tm == null) {
+                Log.e(TAG, "TelephonyManager is null");
+                return false;
             }
 
-            if (callForwardingInfo != null
-                    && callForwardingInfo[i] != null
-                    && subId != SubscriptionManager.INVALID_SUBSCRIPTION_ID) {
-                Log.d(TAG, "Restore call forwarding to " + callForwardingInfo[i]);
-                tm.createForSubscriptionId(subId)
-                        .setCallForwarding(callForwardingInfo[i], null, null);
+            if (callWaitingStatus == null || callForwardingInfo == null) {
+                Log.e(TAG, "CallWaitingStatus or CallForwardingInfo array is null");
+                return false;
             }
+
+            int slotCount = tm.getActiveModemCount();
+            if (callWaitingStatus.length != slotCount || callForwardingInfo.length != slotCount) {
+                Log.e(TAG, "The length of CallWaitingStatus and CallForwardingInfo array"
+                        + " should be the same as phone count.");
+                return false;
+            }
+
+            Executor executor = Executors.newSingleThreadExecutor();
+
+            for (int i = 0; i < slotCount; i++) {
+                int subId = getSubId(i);
+                if (subId == SubscriptionManager.INVALID_SUBSCRIPTION_ID) {
+                    continue;
+                }
+
+                mSteps.add(new RestoreCallWaitingCommand(
+                        tm, executor, callWaitingStatus[i], subId));
+
+                if (callForwardingInfo[i] != null) {
+                    mSteps.add(new RestoreCallForwardingCommand(
+                            tm, executor, callForwardingInfo[i], subId));
+                }
+            }
+
+            return true;
+        }
+
+        /* package */ void startProcess() {
+            int index = 0;
+
+            while (index < mSteps.size()) {
+                UpdateCommand currentStep = mSteps.get(index);
+                Log.d(TAG, "processing : " + currentStep);
+
+                try {
+                    currentStep.process();
+                    index++;
+                } catch (Exception e) {
+                    Log.e(TAG, "Failed on : " + currentStep, e);
+                }
+            }
+        }
+    }
+
+    class RestoreCallForwardingCommand extends UpdateCommand<Integer> {
+        private SettableFuture<Boolean> mResultFuture = SettableFuture.create();
+        private CallForwardingInfo mCallForwardingInfo;
+
+        /* package */ RestoreCallForwardingCommand(TelephonyManager tm, Executor executor,
+                CallForwardingInfo mCallForwardingInfo, int subId) {
+            super(tm, executor, subId);
+            this.mCallForwardingInfo = mCallForwardingInfo;
+        }
+
+        @Override
+        public boolean process() throws Exception {
+            Log.d(TAG, "Restore call forwarding to " + mCallForwardingInfo);
+            tm.createForSubscriptionId(subId).setCallForwarding(mCallForwardingInfo, executor,
+                    this::updateStatusCallBack);
+            return mResultFuture.get();
+        }
+
+        @Override
+        void onRestore() {
+        }
+
+        private void updateStatusCallBack(int result) {
+            Log.d(TAG, "updateStatusCallBack for CallForwarding: " + result);
+            mResultFuture.set(true);
+        }
+    }
+
+    class RestoreCallWaitingCommand extends UpdateCommand<Integer> {
+        private SettableFuture<Boolean> mResultFuture = SettableFuture.create();
+        private boolean mCallWaitingStatus;
+
+        /* package */ RestoreCallWaitingCommand(TelephonyManager tm, Executor executor,
+                boolean mCallWaitingStatus, int subId) {
+            super(tm, executor, subId);
+            this.mCallWaitingStatus = mCallWaitingStatus;
+        }
+
+        @Override
+        public boolean process() throws Exception {
+            Log.d(TAG, "Restore call waiting to " + mCallWaitingStatus);
+            tm.createForSubscriptionId(subId).setCallWaitingEnabled(mCallWaitingStatus, executor,
+                    this::updateStatusCallBack);
+            return mResultFuture.get();
+        }
+
+        @Override
+        void onRestore() {
+        }
+
+        private void updateStatusCallBack(int result) {
+            Log.d(TAG, "updateStatusCallBack for CallWaiting: " + result);
+            mResultFuture.set(true);
         }
     }
 
diff --git a/src/com/android/settings/sim/smartForwarding/SmartForwardingActivity.java b/src/com/android/settings/sim/smartForwarding/SmartForwardingActivity.java
index 95ac999..4f05a8b 100644
--- a/src/com/android/settings/sim/smartForwarding/SmartForwardingActivity.java
+++ b/src/com/android/settings/sim/smartForwarding/SmartForwardingActivity.java
@@ -85,12 +85,7 @@
 
     public void enableSmartForwarding(String[] phoneNumber) {
         // Pop-up ongoing dialog
-        ProgressDialog dialog = new ProgressDialog(this);
-        dialog.setTitle(R.string.smart_forwarding_ongoing_title);
-        dialog.setIndeterminate(true);
-        dialog.setMessage(getText(R.string.smart_forwarding_ongoing_text));
-        dialog.setCancelable(false);
-        dialog.show();
+        ProgressDialog dialog = showOngoingDialog();
 
         // Enable feature
         ListenableFuture<FeatureResult> enableTask =
@@ -140,6 +135,9 @@
         boolean[] callWaitingStatus = getAllSlotCallWaitingStatus(this, tm);
         CallForwardingInfo[] callForwardingInfo = getAllSlotCallForwardingStatus(this, sm, tm);
 
+        // Pop-up ongoing dialog
+        ProgressDialog dialog = showOngoingDialog();
+
         // Disable feature
         ListenableFuture disableTask = service.submit(new DisableSmartForwardingTask(
                 tm, callWaitingStatus, callForwardingInfo));
@@ -147,11 +145,13 @@
             @Override
             public void onSuccess(Object result) {
                 clearAllBackupData(SmartForwardingActivity.this, sm, tm);
+                dialog.dismiss();
             }
 
             @Override
             public void onFailure(Throwable t) {
                 Log.e(TAG, "Disable Feature exception" + t);
+                dialog.dismiss();
             }
         }, ContextCompat.getMainExecutor(this));
     }
@@ -174,4 +174,15 @@
                 .create();
         mDialog.show();
     }
+
+    private ProgressDialog showOngoingDialog() {
+        ProgressDialog dialog = new ProgressDialog(this);
+        dialog.setTitle(R.string.smart_forwarding_ongoing_title);
+        dialog.setIndeterminate(true);
+        dialog.setMessage(getText(R.string.smart_forwarding_ongoing_text));
+        dialog.setCancelable(false);
+        dialog.show();
+
+        return dialog;
+    }
 }
diff --git a/src/com/android/settings/users/OWNERS b/src/com/android/settings/users/OWNERS
new file mode 100644
index 0000000..6afd0ad
--- /dev/null
+++ b/src/com/android/settings/users/OWNERS
@@ -0,0 +1 @@
+include platform/frameworks/base:/MULTIUSER_OWNERS
diff --git a/src/com/android/settings/wfd/WifiDisplayPreferenceController.java b/src/com/android/settings/wfd/WifiDisplayPreferenceController.java
index b15396b..bef16f8 100644
--- a/src/com/android/settings/wfd/WifiDisplayPreferenceController.java
+++ b/src/com/android/settings/wfd/WifiDisplayPreferenceController.java
@@ -65,7 +65,6 @@
     public WifiDisplayPreferenceController(Context context, String key) {
         super(context, key);
         mRouter = context.getSystemService(MediaRouter.class);
-        mRouter.setRouterGroupId(MediaRouter.MIRRORING_GROUP_ID);
     }
 
     @Override
diff --git a/tests/Enable16KbTests/src/com/android/test/Enable16KbTest.java b/tests/Enable16KbTests/src/com/android/test/Enable16KbTest.java
index b611d61..e4ebdef 100644
--- a/tests/Enable16KbTests/src/com/android/test/Enable16KbTest.java
+++ b/tests/Enable16KbTests/src/com/android/test/Enable16KbTest.java
@@ -76,11 +76,13 @@
         installTestApp();
 
         // Enable developer option and switch to 16kb kernel and Check page size
+        getDevice().enableAdbRoot();
         runTestAndWait(SWITCH_TO_16KB);
         result = getDevice().executeShellCommand("getconf PAGE_SIZE");
         assertEquals("16384", result.strip());
 
         // switch back to 4kb kernel and check page size
+        getDevice().enableAdbRoot();
         runTestAndWait(SWITCH_TO_4KB);
         result = getDevice().executeShellCommand("getconf PAGE_SIZE");
         assertEquals("4096", result.strip());
diff --git a/tests/Enable16KbTests/test_16kb_app/src/com/android/settings/development/test/Enable16KbDeviceTest.java b/tests/Enable16KbTests/test_16kb_app/src/com/android/settings/development/test/Enable16KbDeviceTest.java
index e5ccbb9..9e489cf 100644
--- a/tests/Enable16KbTests/test_16kb_app/src/com/android/settings/development/test/Enable16KbDeviceTest.java
+++ b/tests/Enable16KbTests/test_16kb_app/src/com/android/settings/development/test/Enable16KbDeviceTest.java
@@ -68,6 +68,8 @@
             throw new RuntimeException("failed to freeze device orientation", e);
         }
 
+        mDevice.executeShellCommand("am start -a com.android.setupwizard.FOUR_CORNER_EXIT");
+        mDevice.waitForWindowUpdate(null, TIMEOUT);
         mDevice.executeShellCommand("input keyevent KEYCODE_WAKEUP");
         mDevice.executeShellCommand("wm dismiss-keyguard");
     }
@@ -150,6 +152,7 @@
 
     private void openPersistentNotification(String title) {
         mDevice.openNotification();
+        mDevice.waitForWindowUpdate(null, TIMEOUT);
         verifyTextOnScreen(title);
         mDevice.wait(Until.findObject(By.text(title)), TIMEOUT).click();
         mDevice.waitForWindowUpdate(null, TIMEOUT);
diff --git a/tests/robotests/src/com/android/settings/development/OWNERS b/tests/robotests/src/com/android/settings/development/OWNERS
new file mode 100644
index 0000000..92a75e2
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/development/OWNERS
@@ -0,0 +1,2 @@
+# DesktopModePreferenceControllerTest
+per-file DesktopModePreferenceControllerTest.java=file:platform/frameworks/base:/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/OWNERS
diff --git a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppBatteryPreferenceTest.kt b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppBatteryPreferenceTest.kt
index b7b96b0..a724a8c 100644
--- a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppBatteryPreferenceTest.kt
+++ b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppBatteryPreferenceTest.kt
@@ -20,13 +20,14 @@
 import android.content.pm.ApplicationInfo
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.ui.platform.LocalContext
+import androidx.compose.ui.test.ExperimentalTestApi
 import androidx.compose.ui.test.assertIsDisplayed
-import androidx.compose.ui.test.assertIsEnabled
 import androidx.compose.ui.test.assertIsNotDisplayed
 import androidx.compose.ui.test.assertIsNotEnabled
+import androidx.compose.ui.test.hasText
 import androidx.compose.ui.test.hasTextExactly
+import androidx.compose.ui.test.isEnabled
 import androidx.compose.ui.test.junit4.createComposeRule
-import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.onRoot
 import androidx.compose.ui.test.performClick
 import androidx.test.core.app.ApplicationProvider
@@ -36,18 +37,23 @@
 import com.android.settings.fuelgauge.AdvancedPowerUsageDetail
 import com.android.settings.fuelgauge.batteryusage.BatteryChartPreferenceController
 import com.android.settings.fuelgauge.batteryusage.BatteryDiffEntry
+import com.android.settingslib.spa.testutils.delay
 import com.android.settingslib.spaprivileged.model.app.userId
+import kotlinx.coroutines.runBlocking
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.mock
 import org.mockito.MockitoSession
-import org.mockito.Spy
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
+import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
-import org.mockito.Mockito.`when` as whenever
 
+@OptIn(ExperimentalTestApi::class)
 @RunWith(AndroidJUnit4::class)
 class AppBatteryPreferenceTest {
     @get:Rule
@@ -55,29 +61,28 @@
 
     private lateinit var mockSession: MockitoSession
 
-    @Spy
-    private val context: Context = ApplicationProvider.getApplicationContext()
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {}
 
-    @Spy
-    private val resources = context.resources
+    private val resources = spy(context.resources) {
+        on { getBoolean(R.bool.config_show_app_info_settings_battery) } doReturn true
+    }
 
     @Before
     fun setUp() {
         mockSession = ExtendedMockito.mockitoSession()
-            .initMocks(this)
             .mockStatic(BatteryChartPreferenceController::class.java)
             .mockStatic(AdvancedPowerUsageDetail::class.java)
             .strictness(Strictness.LENIENT)
             .startMocking()
         whenever(context.resources).thenReturn(resources)
-        whenever(resources.getBoolean(R.bool.config_show_app_info_settings_battery))
-            .thenReturn(true)
     }
 
     private fun mockBatteryDiffEntry(batteryDiffEntry: BatteryDiffEntry?) {
-        whenever(BatteryChartPreferenceController.getAppBatteryUsageData(
-            context, PACKAGE_NAME, APP.userId
-        )).thenReturn(batteryDiffEntry)
+        whenever(
+            BatteryChartPreferenceController.getAppBatteryUsageData(
+                context, PACKAGE_NAME, APP.userId
+            )
+        ).thenReturn(batteryDiffEntry)
     }
 
     @After
@@ -87,8 +92,9 @@
 
     @Test
     fun whenConfigIsFalse_notDisplayed() {
-        whenever(resources.getBoolean(R.bool.config_show_app_info_settings_battery))
-            .thenReturn(false)
+        resources.stub {
+            on { getBoolean(R.bool.config_show_app_info_settings_battery) } doReturn false
+        }
 
         setContent()
 
@@ -112,47 +118,47 @@
 
         setContent()
 
-        composeTestRule.onNode(
+        composeTestRule.waitUntilExactlyOneExists(
             hasTextExactly(
                 context.getString(R.string.battery_details_title),
                 context.getString(R.string.no_battery_summary),
-            ),
-        ).assertIsDisplayed().assertIsEnabled()
+            ) and isEnabled(),
+        )
     }
 
     @Test
     fun noConsumePower() {
-        val batteryDiffEntry = mock(BatteryDiffEntry::class.java).apply {
-            mConsumePower = 0.0
-        }
+        val batteryDiffEntry = mock<BatteryDiffEntry>().apply { mConsumePower = 0.0 }
         mockBatteryDiffEntry(batteryDiffEntry)
 
         setContent()
 
-        composeTestRule.onNodeWithText(context.getString(R.string.no_battery_summary))
-            .assertIsDisplayed()
+        composeTestRule.waitUntilExactlyOneExists(
+            hasText(context.getString(R.string.no_battery_summary))
+        )
     }
 
     @Test
     fun hasConsumePower() {
-        val batteryDiffEntry = mock(BatteryDiffEntry::class.java).apply {
-            mConsumePower = 12.3
-        }
-        whenever(batteryDiffEntry.percentage).thenReturn(45.6)
+        val batteryDiffEntry = mock<BatteryDiffEntry> {
+            on { percentage } doReturn 45.6
+        }.apply { mConsumePower = 12.3 }
         mockBatteryDiffEntry(batteryDiffEntry)
 
         setContent()
 
-        composeTestRule.onNodeWithText("46% use since last full charge").assertIsDisplayed()
+        composeTestRule.waitUntilExactlyOneExists(hasText("46% use since last full charge"))
     }
 
     @Test
     fun whenClick_openDetailsPage() {
-        val batteryDiffEntry = mock(BatteryDiffEntry::class.java)
-        whenever(batteryDiffEntry.percentage).thenReturn(10.0)
+        val batteryDiffEntry = mock<BatteryDiffEntry> {
+            on { percentage } doReturn 10.0
+        }.apply { mConsumePower = 12.3 }
         mockBatteryDiffEntry(batteryDiffEntry)
 
         setContent()
+        composeTestRule.waitUntilExactlyOneExists(hasText("10% use since last full charge"))
         composeTestRule.onRoot().performClick()
 
         ExtendedMockito.verify {
@@ -178,7 +184,7 @@
     }
 
     private companion object {
-        const val PACKAGE_NAME = "packageName"
+        const val PACKAGE_NAME = "package.name"
         const val UID = 123
         val APP = ApplicationInfo().apply {
             packageName = PACKAGE_NAME