diff --git a/src/com/android/settings/fuelgauge/batteryusage/BatteryCallbackData.java b/src/com/android/settings/fuelgauge/batteryusage/BatteryCallbackData.java
deleted file mode 100644
index 17e5f5c..0000000
--- a/src/com/android/settings/fuelgauge/batteryusage/BatteryCallbackData.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Copyright (C) 2022 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.settings.fuelgauge.batteryusage;
-
-import androidx.annotation.Nullable;
-
-import java.util.Locale;
-import java.util.Map;
-
-/** Wraps the battery usage data and device screen-on time data used for battery usage page. */
-public class BatteryCallbackData {
-
-    // The usage app data used for rendering app list.
-    private final Map<Integer, Map<Integer, BatteryDiffData>> mBatteryUsageMap;
-    // The device screen-on time data.
-    private final Map<Integer, Map<Integer, Long>> mDeviceScreenOnTime;
-
-    public BatteryCallbackData(
-            @Nullable Map<Integer, Map<Integer, BatteryDiffData>> batteryUsageMap,
-            @Nullable Map<Integer, Map<Integer, Long>> deviceScreenOnTime) {
-        mBatteryUsageMap = batteryUsageMap;
-        mDeviceScreenOnTime = deviceScreenOnTime;
-    }
-
-    public Map<Integer, Map<Integer, BatteryDiffData>> getBatteryUsageMap() {
-        return mBatteryUsageMap;
-    }
-
-    public Map<Integer, Map<Integer, Long>> getDeviceScreenOnTime() {
-        return mDeviceScreenOnTime;
-    }
-
-    @Override
-    public String toString() {
-        return String.format(Locale.ENGLISH,
-                "batteryUsageMap: %s; deviceScreenOnTime: %s",
-                mBatteryUsageMap,
-                mDeviceScreenOnTime);
-    }
-}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/BatteryChartPreferenceController.java b/src/com/android/settings/fuelgauge/batteryusage/BatteryChartPreferenceController.java
index 67d3969..17d9c8a 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/BatteryChartPreferenceController.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/BatteryChartPreferenceController.java
@@ -110,8 +110,6 @@
     int mHourlyChartIndex = BatteryChartViewModel.SELECTED_INDEX_ALL;
     @VisibleForTesting
     Map<Integer, Map<Integer, BatteryDiffData>> mBatteryUsageMap;
-    @VisibleForTesting
-    Map<Integer, Map<Integer, Long>> mScreenOnTimeMap;
 
     private boolean mIs24HourFormat;
     private boolean mHourlyChartVisible = true;
@@ -219,9 +217,8 @@
         animateBatteryChartViewGroup();
         final BatteryLevelData batteryLevelData =
                 DataProcessManager.getBatteryLevelData(mContext, mHandler, batteryHistoryMap,
-                        batteryCallbackData -> {
-                            mBatteryUsageMap = batteryCallbackData.getBatteryUsageMap();
-                            mScreenOnTimeMap = batteryCallbackData.getDeviceScreenOnTime();
+                        batteryUsageMap -> {
+                            mBatteryUsageMap = batteryUsageMap;
                             logScreenUsageTime();
                             refreshUi();
                         });
@@ -336,16 +333,15 @@
             return false;
         }
 
-        if (mOnScreenOnTimeUpdatedListener != null && mScreenOnTimeMap != null
-                && mScreenOnTimeMap.get(mDailyChartIndex) != null) {
-            mOnScreenOnTimeUpdatedListener.onScreenOnTimeUpdated(
-                    mScreenOnTimeMap.get(mDailyChartIndex).get(mHourlyChartIndex),
-                    getSlotInformation());
-        }
         if (mOnBatteryUsageUpdatedListener != null && mBatteryUsageMap != null
                 && mBatteryUsageMap.get(mDailyChartIndex) != null) {
             final BatteryDiffData slotUsageData =
                     mBatteryUsageMap.get(mDailyChartIndex).get(mHourlyChartIndex);
+            if (slotUsageData != null) {
+                mOnScreenOnTimeUpdatedListener.onScreenOnTimeUpdated(
+                        slotUsageData.getScreenOnTime(),
+                        getSlotInformation());
+            }
             mOnBatteryUsageUpdatedListener.onBatteryUsageUpdated(
                     slotUsageData, getSlotInformation(), isBatteryUsageMapNullOrEmpty());
         }
@@ -503,17 +499,19 @@
     }
 
     private void logScreenUsageTime() {
-        if (mBatteryUsageMap == null || mScreenOnTimeMap == null) {
+        if (mBatteryUsageMap == null) {
             return;
         }
-        final long totalScreenOnTime =
-                mScreenOnTimeMap
-                        .get(BatteryChartViewModel.SELECTED_INDEX_ALL)
-                        .get(BatteryChartViewModel.SELECTED_INDEX_ALL);
+        final BatteryDiffData allBatteryDiffData = mBatteryUsageMap.get(
+                BatteryChartViewModel.SELECTED_INDEX_ALL).get(
+                BatteryChartViewModel.SELECTED_INDEX_ALL);
+        if (allBatteryDiffData == null) {
+            return;
+        }
         mMetricsFeatureProvider.action(
                 mPrefContext,
                 SettingsEnums.ACTION_BATTERY_USAGE_SCREEN_ON_TIME,
-                (int) totalScreenOnTime);
+                (int) allBatteryDiffData.getScreenOnTime());
         mMetricsFeatureProvider.action(
                 mPrefContext,
                 SettingsEnums.ACTION_BATTERY_USAGE_FOREGROUND_USAGE_TIME,
diff --git a/src/com/android/settings/fuelgauge/batteryusage/BatteryDiffData.java b/src/com/android/settings/fuelgauge/batteryusage/BatteryDiffData.java
index 37eddf3..47ae568 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/BatteryDiffData.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/BatteryDiffData.java
@@ -34,17 +34,20 @@
 public class BatteryDiffData {
     static final double SMALL_PERCENTAGE_THRESHOLD = 1f;
 
+    private final long mScreenOnTime;
     private final List<BatteryDiffEntry> mAppEntries;
     private final List<BatteryDiffEntry> mSystemEntries;
 
     /** Constructor for the diff entries. */
     public BatteryDiffData(
             final Context context,
+            final long screenOnTime,
             final @NonNull List<BatteryDiffEntry> appDiffEntries,
             final @NonNull List<BatteryDiffEntry> systemDiffEntries,
             final @NonNull Set<String> systemAppsPackageNames,
             final @NonNull Set<Integer> systemAppsUids,
             final boolean isAccumulated) {
+        mScreenOnTime = screenOnTime;
         mAppEntries = appDiffEntries;
         mSystemEntries = systemDiffEntries;
 
@@ -60,6 +63,10 @@
         processAndSortEntries(mSystemEntries);
     }
 
+    public long getScreenOnTime() {
+        return mScreenOnTime;
+    }
+
     public List<BatteryDiffEntry> getAppDiffEntryList() {
         return mAppEntries;
     }
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
index 6b007de..0f67e6a 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
@@ -353,9 +353,9 @@
     }
 
     private void loadAndApplyBatteryMapFromServiceOnly() {
-        new AsyncTask<Void, Void, BatteryCallbackData>() {
+        new AsyncTask<Void, Void, Map<Integer, Map<Integer, BatteryDiffData>>>() {
             @Override
-            protected BatteryCallbackData doInBackground(Void... voids) {
+            protected Map<Integer, Map<Integer, BatteryDiffData>> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
                 final Map<Integer, Map<Integer, BatteryDiffData>> batteryUsageMap =
                         DataProcessor.getBatteryUsageMapFromStatsService(mContext);
@@ -363,18 +363,18 @@
                 Log.d(TAG, String.format(
                         "execute loadAndApplyBatteryMapFromServiceOnly size=%d in %d/ms",
                         batteryUsageMap.size(), (System.currentTimeMillis() - startTime)));
-                return new BatteryCallbackData(batteryUsageMap, /*deviceScreenOnTime=*/ null);
+                return batteryUsageMap;
             }
 
             @Override
             protected void onPostExecute(
-                    final BatteryCallbackData batteryCallbackData) {
+                    final Map<Integer, Map<Integer, BatteryDiffData>> batteryUsageMap) {
                 // Set the unused variables to null.
                 mContext = null;
                 // Post results back to main thread to refresh UI.
                 if (mHandler != null && mCallbackFunction != null) {
                     mHandler.post(() -> {
-                        mCallbackFunction.onBatteryCallbackDataLoaded(batteryCallbackData);
+                        mCallbackFunction.onBatteryCallbackDataLoaded(batteryUsageMap);
                     });
                 }
             }
@@ -413,24 +413,23 @@
     }
 
     private void generateFinalDataAndApplyCallback() {
-        new AsyncTask<Void, Void, BatteryCallbackData>() {
+        new AsyncTask<Void, Void, Map<Integer, Map<Integer, BatteryDiffData>>>() {
             @Override
-            protected BatteryCallbackData doInBackground(Void... voids) {
+            protected Map<Integer, Map<Integer, BatteryDiffData>> doInBackground(Void... voids) {
                 final long startTime = System.currentTimeMillis();
                 final Map<Integer, Map<Integer, BatteryDiffData>> batteryUsageMap =
                         DataProcessor.getBatteryUsageMap(
                                 mContext, mHourlyBatteryLevelsPerDay, mBatteryHistoryMap,
                                 mAppUsagePeriodMap);
-                final Map<Integer, Map<Integer, Long>> deviceScreenOnTime =
-                        DataProcessor.getDeviceScreenOnTime(mAppUsagePeriodMap);
                 DataProcessor.loadLabelAndIcon(batteryUsageMap);
                 Log.d(TAG, String.format("execute generateFinalDataAndApplyCallback in %d/ms",
                         (System.currentTimeMillis() - startTime)));
-                return new BatteryCallbackData(batteryUsageMap, deviceScreenOnTime);
+                return batteryUsageMap;
             }
 
             @Override
-            protected void onPostExecute(final BatteryCallbackData batteryCallbackData) {
+            protected void onPostExecute(
+                    final Map<Integer, Map<Integer, BatteryDiffData>> batteryUsageMap) {
                 // Set the unused variables to null.
                 mContext = null;
                 mHourlyBatteryLevelsPerDay = null;
@@ -438,7 +437,7 @@
                 // Post results back to main thread to refresh UI.
                 if (mHandler != null && mCallbackFunction != null) {
                     mHandler.post(() -> {
-                        mCallbackFunction.onBatteryCallbackDataLoaded(batteryCallbackData);
+                        mCallbackFunction.onBatteryCallbackDataLoaded(batteryUsageMap);
                     });
                 }
             }
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
index 9aee2f7..78541c1 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessor.java
@@ -122,7 +122,8 @@
     /** A callback listener when battery usage loading async task is executed. */
     public interface UsageMapAsyncResponse {
         /** The callback function when batteryUsageMap is loaded. */
-        void onBatteryCallbackDataLoaded(BatteryCallbackData batteryCallbackData);
+        void onBatteryCallbackDataLoaded(
+                Map<Integer, Map<Integer, BatteryDiffData>> batteryCallbackData);
     }
 
     private DataProcessor() {
@@ -364,35 +365,6 @@
     }
 
     /**
-     * @return Returns the device screen-on time data.
-     *
-     * <p>There could be 2 cases of the returned value:</p>
-     * <ul>
-     * <li>null: empty or invalid data.</li>
-     * <li>non-null: must be a 2d map and composed by 3 parts:</li>
-     * <p>  1 - [SELECTED_INDEX_ALL][SELECTED_INDEX_ALL]</p>
-     * <p>  2 - [0][SELECTED_INDEX_ALL] ~ [maxDailyIndex][SELECTED_INDEX_ALL]</p>
-     * <p>  3 - [0][0] ~ [maxDailyIndex][maxHourlyIndex]</p>
-     * </ul>
-     *
-     * <p>The structure is consistent with the battery usage map returned by
-     * {@code getBatteryUsageMap}.</p>
-     */
-    @Nullable
-    public static Map<Integer, Map<Integer, Long>> getDeviceScreenOnTime(
-            final Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>>
-                    appUsagePeriodMap) {
-        if (appUsagePeriodMap == null) {
-            return null;
-        }
-        final Map<Integer, Map<Integer, Long>> deviceScreenOnTime = new ArrayMap<>();
-        insertHourlyDeviceScreenOnTime(appUsagePeriodMap, deviceScreenOnTime);
-        insertDailyDeviceScreenOnTime(appUsagePeriodMap, deviceScreenOnTime);
-        insertAllDeviceScreenOnTime(deviceScreenOnTime);
-        return deviceScreenOnTime;
-    }
-
-    /**
      * Generates the list of {@link BatteryEntry} from the supplied {@link BatteryUsageStats}.
      */
     @Nullable
@@ -678,7 +650,7 @@
 
         final Set<String> systemAppsPackageNames = getSystemAppsPackageNames(context);
         final Set<Integer> systemAppsUids = getSystemAppsUids(context);
-        return new BatteryDiffData(context, appEntries, systemEntries,
+        return new BatteryDiffData(context, /* screenOnTime= */ 0L, appEntries, systemEntries,
                 systemAppsPackageNames, systemAppsUids, /* isAccumulated= */ false);
     }
 
@@ -988,84 +960,6 @@
                 eventTime);
     }
 
-    private static void insertHourlyDeviceScreenOnTime(
-            final Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>>
-                    appUsagePeriodMap,
-            final Map<Integer, Map<Integer, Long>> resultMap) {
-        for (final int dailyIndex : appUsagePeriodMap.keySet()) {
-            final Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>> dailyAppUsageMap =
-                    appUsagePeriodMap.get(dailyIndex);
-            final Map<Integer, Long> dailyScreenOnTime = new ArrayMap<>();
-            resultMap.put(dailyIndex, dailyScreenOnTime);
-            if (dailyAppUsageMap == null) {
-                continue;
-            }
-
-            for (final int hourlyIndex : dailyAppUsageMap.keySet()) {
-                final Map<Long, Map<String, List<AppUsagePeriod>>> appUsageMap =
-                        dailyAppUsageMap.get(hourlyIndex);
-                if (appUsageMap == null || appUsageMap.isEmpty()) {
-                    dailyScreenOnTime.put(hourlyIndex, 0L);
-                } else {
-                    final List<AppUsagePeriod> flatUsageList = new ArrayList<>();
-                    for (final long userId: appUsageMap.keySet()) {
-                        if (appUsageMap.get(userId) == null) {
-                            continue;
-                        }
-                        for (final String packageName: appUsageMap.get(userId).keySet()) {
-                            final List<AppUsagePeriod> appUsagePeriodList =
-                                    appUsageMap.get(userId).get(packageName);
-                            if (appUsagePeriodList != null && !appUsagePeriodList.isEmpty()) {
-                                flatUsageList.addAll(appUsagePeriodList);
-                            }
-                        }
-                    }
-                    // Compute the screen on time and make sure it won't exceed the threshold.
-                    final long screenOnTime = Math.min(
-                            (long) TOTAL_HOURLY_TIME_THRESHOLD, getScreenOnTime(flatUsageList));
-                    dailyScreenOnTime.put(hourlyIndex, screenOnTime);
-                }
-            }
-        }
-    }
-
-    private static void insertDailyDeviceScreenOnTime(
-            final Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>>
-                    appUsagePeriodMap,
-            final Map<Integer, Map<Integer, Long>> resultMap) {
-        for (final int dailyIndex : appUsagePeriodMap.keySet()) {
-            Map<Integer, Long> dailyResultMap = resultMap.get(dailyIndex);
-            if (dailyResultMap == null) {
-                dailyResultMap = new ArrayMap<>();
-                resultMap.put(dailyIndex, dailyResultMap);
-            }
-            dailyResultMap.put(
-                    SELECTED_INDEX_ALL,
-                    getAccumulatedScreenOnTime(dailyResultMap));
-        }
-    }
-
-    private static void insertAllDeviceScreenOnTime(
-            final Map<Integer, Map<Integer, Long>> resultMap) {
-        final Map<Integer, Long> dailyAllMap = new ArrayMap<>();
-        resultMap.keySet().forEach(
-                key -> dailyAllMap.put(key, resultMap.get(key).get(SELECTED_INDEX_ALL)));
-        final Map<Integer, Long> allUsageMap = new ArrayMap<>();
-        allUsageMap.put(SELECTED_INDEX_ALL, getAccumulatedScreenOnTime(dailyAllMap));
-        resultMap.put(SELECTED_INDEX_ALL, allUsageMap);
-    }
-
-    private static long getAccumulatedScreenOnTime(final Map<Integer, Long> screenOnTimeMap) {
-        if (screenOnTimeMap == null || screenOnTimeMap.isEmpty()) {
-            return 0;
-        }
-        long sum = 0;
-        for (final int index : screenOnTimeMap.keySet()) {
-            sum += screenOnTimeMap.get(index) == null ? 0 : screenOnTimeMap.get(index);
-        }
-        return sum;
-    }
-
     @Nullable
     private static UsageEvents getAppUsageEventsForUser(
             Context context, final UserManager userManager, final int userID,
@@ -1113,7 +1007,7 @@
             batteryHistEntryList = convertToBatteryHistEntry(batteryEntryList, batteryUsageStats);
             closeBatteryUsageStats(batteryUsageStats);
         } catch (RuntimeException e) {
-            Log.e(TAG, "load batteryUsageStats:" + e);
+            Log.e(TAG, "load batteryUsageStats:", e);
         }
 
         return batteryHistEntryList;
@@ -1406,6 +1300,26 @@
             final Set<Integer> systemAppsUids,
             final Map<Long, Map<String, List<AppUsagePeriod>>> appUsageMap,
             final List<Map<String, BatteryHistEntry>> slotBatteryHistoryList) {
+        long slotScreenOnTime = 0L;
+        if (appUsageMap != null) {
+            final List<AppUsagePeriod> flatAppUsagePeriodList = new ArrayList<>();
+            for (final long userId : appUsageMap.keySet()) {
+                if ((userId != currentUserId && userId != workProfileUserId)
+                        || appUsageMap.get(userId) == null) {
+                    continue;
+                }
+                for (final String packageName : appUsageMap.get(userId).keySet()) {
+                    final List<AppUsagePeriod> appUsagePeriodList =
+                            appUsageMap.get(userId).get(packageName);
+                    if (appUsagePeriodList != null) {
+                        flatAppUsagePeriodList.addAll(appUsagePeriodList);
+                    }
+                }
+            }
+            slotScreenOnTime =
+                    Math.min(slotDuration, getScreenOnTime(flatAppUsagePeriodList));
+        }
+
         final List<BatteryDiffEntry> appEntries = new ArrayList<>();
         final List<BatteryDiffEntry> systemEntries = new ArrayList<>();
 
@@ -1550,7 +1464,7 @@
             return null;
         }
 
-        return new BatteryDiffData(context, appEntries, systemEntries,
+        return new BatteryDiffData(context, slotScreenOnTime, appEntries, systemEntries,
                 systemAppsPackageNames, systemAppsUids, /* isAccumulated= */ false);
     }
 
@@ -1608,19 +1522,21 @@
 
     @Nullable
     private static BatteryDiffData getAccumulatedUsageDiffData(
-            final Context context, final Collection<BatteryDiffData> diffEntryListData) {
+            final Context context, final Collection<BatteryDiffData> batteryDiffDataList) {
         final Map<String, BatteryDiffEntry> diffEntryMap = new ArrayMap<>();
         final List<BatteryDiffEntry> appEntries = new ArrayList<>();
         final List<BatteryDiffEntry> systemEntries = new ArrayList<>();
 
-        for (BatteryDiffData diffEntryList : diffEntryListData) {
-            if (diffEntryList == null) {
+        long totalScreenOnTime = 0;
+        for (BatteryDiffData batteryDiffData : batteryDiffDataList) {
+            if (batteryDiffData == null) {
                 continue;
             }
-            for (BatteryDiffEntry entry : diffEntryList.getAppDiffEntryList()) {
+            totalScreenOnTime += batteryDiffData.getScreenOnTime();
+            for (BatteryDiffEntry entry : batteryDiffData.getAppDiffEntryList()) {
                 computeUsageDiffDataPerEntry(entry, diffEntryMap);
             }
-            for (BatteryDiffEntry entry : diffEntryList.getSystemDiffEntryList()) {
+            for (BatteryDiffEntry entry : batteryDiffData.getSystemDiffEntryList()) {
                 computeUsageDiffDataPerEntry(entry, diffEntryMap);
             }
         }
@@ -1634,8 +1550,8 @@
             }
         }
 
-        return diffEntryList.isEmpty() ? null : new BatteryDiffData(context, appEntries,
-                systemEntries, /* systemAppsPackageNames= */ new ArraySet<>(),
+        return diffEntryList.isEmpty() ? null : new BatteryDiffData(context, totalScreenOnTime,
+                appEntries, systemEntries, /* systemAppsPackageNames= */ new ArraySet<>(),
                 /* systemAppsUids= */ new ArraySet<>(), /* isAccumulated= */ true);
     }
 
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageBreakdownControllerTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageBreakdownControllerTest.java
index 8d0a839..3a9ce2b 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageBreakdownControllerTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/BatteryUsageBreakdownControllerTest.java
@@ -107,8 +107,9 @@
                 mBatteryHistEntry);
         mBatteryDiffEntry = spy(mBatteryDiffEntry);
         mBatteryUsageBreakdownController.mBatteryDiffData =
-                new BatteryDiffData(mContext, Arrays.asList(mBatteryDiffEntry), Arrays.asList(),
-                        Set.of(), Set.of(), /* isAccumulated= */ false);
+                new BatteryDiffData(mContext, /* screenOnTime= */ 0L,
+                        Arrays.asList(mBatteryDiffEntry), Arrays.asList(), Set.of(), Set.of(),
+                        /* isAccumulated= */ false);
         // Adds fake testing data.
         BatteryDiffEntry.sResourceCache.put(
                 "fakeBatteryDiffEntryKey",
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
index ea2db86..d69eec7 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/DataProcessorTest.java
@@ -326,62 +326,6 @@
     }
 
     @Test
-    public void getDeviceScreenOnTime_returnExpectedResult() {
-        final Map<Integer, Map<Integer, Map<Long, Map<String, List<AppUsagePeriod>>>>>
-                appUsagePeriodMap = new HashMap<>();
-        appUsagePeriodMap.put(0, new HashMap<>());
-        appUsagePeriodMap.put(1, new HashMap<>());
-        appUsagePeriodMap.put(2, null);
-        final long userId1 = 1;
-        final long userId2 = 2;
-        // Adds the index [0][0].
-        Map<Long, Map<String, List<AppUsagePeriod>>> appUsageMap = new HashMap<>();
-        Map<String, List<AppUsagePeriod>> userPeriodMap = new HashMap<>();
-        appUsageMap.put(userId1, userPeriodMap);
-        userPeriodMap.put(
-                "package1", List.of(buildAppUsagePeriod(0, 5), buildAppUsagePeriod(5, 7)));
-        userPeriodMap.put("package2", List.of(buildAppUsagePeriod(10, 25)));
-        userPeriodMap = new HashMap<>();
-        appUsageMap.put(userId2, userPeriodMap);
-        userPeriodMap.put("package3", List.of(buildAppUsagePeriod(15, 45)));
-        appUsagePeriodMap.get(0).put(0, appUsageMap);
-        // Adds the index [0][1].
-        appUsageMap = new HashMap<>();
-        userPeriodMap = new HashMap<>();
-        appUsageMap.put(userId1, userPeriodMap);
-        userPeriodMap.put(
-                "package1", List.of(buildAppUsagePeriod(50, 60), buildAppUsagePeriod(70, 80)));
-        appUsagePeriodMap.get(0).put(1, appUsageMap);
-        // Adds the index [1][0].
-        appUsageMap = new HashMap<>();
-        userPeriodMap = new HashMap<>();
-        appUsageMap.put(userId1, userPeriodMap);
-        userPeriodMap.put("package2", List.of(buildAppUsagePeriod(0, 8000000L)));
-        userPeriodMap.put("package3",
-                List.of(buildAppUsagePeriod(10, 15), buildAppUsagePeriod(25, 29)));
-        appUsagePeriodMap.get(1).put(0, appUsageMap);
-
-        final Map<Integer, Map<Integer, Long>> deviceScreenOnTime =
-                DataProcessor.getDeviceScreenOnTime(appUsagePeriodMap);
-
-        assertThat(deviceScreenOnTime.get(0).get(0)).isEqualTo(42);
-        assertThat(deviceScreenOnTime.get(0).get(1)).isEqualTo(20);
-        assertThat(deviceScreenOnTime.get(1).get(0)).isEqualTo(7200000);
-        assertThat(deviceScreenOnTime.get(0).get(DataProcessor.SELECTED_INDEX_ALL)).isEqualTo(62);
-        assertThat(deviceScreenOnTime.get(1).get(DataProcessor.SELECTED_INDEX_ALL))
-                .isEqualTo(7200000);
-        assertThat(deviceScreenOnTime
-                .get(DataProcessor.SELECTED_INDEX_ALL)
-                .get(DataProcessor.SELECTED_INDEX_ALL))
-                .isEqualTo(7200062);
-    }
-
-    @Test
-    public void getDeviceScreenOnTime_nullUsageMap_returnNull() {
-        assertThat(DataProcessor.getDeviceScreenOnTime(null)).isNull();
-    }
-
-    @Test
     public void getHistoryMapWithExpectedTimestamps_emptyHistoryMap_returnEmptyMap() {
         assertThat(DataProcessor
                 .getHistoryMapWithExpectedTimestamps(mContext, new HashMap<>()))
@@ -1128,6 +1072,7 @@
                 resultMap
                         .get(DataProcessor.SELECTED_INDEX_ALL)
                         .get(DataProcessor.SELECTED_INDEX_ALL);
+        assertThat(resultDiffData.getScreenOnTime()).isEqualTo(36L);
         assertBatteryDiffEntry(
                 resultDiffData.getAppDiffEntryList().get(0), currentUserId, /*uid=*/ 2L,
                 ConvertUtils.CONSUMER_TYPE_UID_BATTERY, /*consumePercentage=*/ 50.0,
@@ -1287,6 +1232,7 @@
                 resultMap
                         .get(DataProcessor.SELECTED_INDEX_ALL)
                         .get(DataProcessor.SELECTED_INDEX_ALL);
+        assertThat(resultDiffData.getScreenOnTime()).isEqualTo(0L);
         assertBatteryDiffEntry(
                 resultDiffData.getAppDiffEntryList().get(0), currentUserId, /*uid=*/ 1L,
                 ConvertUtils.CONSUMER_TYPE_UID_BATTERY, /*consumePercentage=*/ 100.0,
@@ -1367,6 +1313,7 @@
                 resultMap
                         .get(DataProcessor.SELECTED_INDEX_ALL)
                         .get(DataProcessor.SELECTED_INDEX_ALL);
+        assertThat(resultDiffData.getScreenOnTime()).isEqualTo(7200000L);
         // Verifies the clipped usage time.
         final float ratio = (float) (7200) / (float) (3600 + 7200);
         final BatteryDiffEntry resultEntry = resultDiffData.getAppDiffEntryList().get(0);