diff --git a/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetector.java b/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetector.java
index 74c43d0..545bb26 100644
--- a/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetector.java
+++ b/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetector.java
@@ -21,7 +21,6 @@
 import android.os.BatteryStats;
 import android.os.SystemClock;
 import android.support.annotation.VisibleForTesting;
-import android.util.ArrayMap;
 
 import com.android.internal.os.BatterySipper;
 import com.android.internal.os.BatteryStatsHelper;
@@ -87,23 +86,10 @@
                 continue;
             }
 
-            final ArrayMap<String, ? extends BatteryStats.Uid.Wakelock> wakelocks =
-                    uid.getWakelockStats();
-            long maxPartialWakeLockMs = 0;
+            final long currentDurationMs = getCurrentDurationMs(uid, rawRealtime);
+            final long backgroundDurationMs = getBackgroundTotalDurationMs(uid, rawRealtime);
 
-            for (int iw = wakelocks.size() - 1; iw >= 0; iw--) {
-                final BatteryStats.Timer timer = wakelocks.valueAt(iw).getWakeTime(
-                        BatteryStats.WAKE_TYPE_PARTIAL);
-                if (timer == null) {
-                    continue;
-                }
-                maxPartialWakeLockMs = Math.max(maxPartialWakeLockMs,
-                        getTotalDurationMs(timer, rawRealtime));
-            }
-
-            // Report application as anomaly if wakelock time is too long
-            // TODO(b/38233034): add more attributes to detect wakelock anomaly
-            if (maxPartialWakeLockMs > mWakeLockThresholdMs) {
+            if (backgroundDurationMs > mWakeLockThresholdMs && currentDurationMs != 0) {
                 final String packageName = mBatteryUtils.getPackageName(uid.getUid());
                 final CharSequence displayName = Utils.getApplicationLabel(mContext,
                         packageName);
@@ -119,16 +105,22 @@
                     anomalies.add(anomaly);
                 }
             }
-
         }
         return anomalies;
     }
 
     @VisibleForTesting
-    long getTotalDurationMs(BatteryStats.Timer timer, long rawRealtime) {
-        if (timer == null) {
-            return 0;
-        }
-        return timer.getTotalDurationMsLocked(rawRealtime);
+    long getCurrentDurationMs(BatteryStats.Uid uid, long elapsedRealtimeMs) {
+        BatteryStats.Timer timer = uid.getAggregatedPartialWakelockTimer();
+
+        return timer != null ? timer.getCurrentDurationMsLocked(elapsedRealtimeMs) : 0;
+    }
+
+    @VisibleForTesting
+    long getBackgroundTotalDurationMs(BatteryStats.Uid uid, long elapsedRealtimeMs) {
+        BatteryStats.Timer timer = uid.getAggregatedPartialWakelockTimer();
+        BatteryStats.Timer subTimer = timer != null ? timer.getSubTimer() : null;
+
+        return subTimer != null ? subTimer.getTotalDurationMsLocked(elapsedRealtimeMs) : 0;
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetectorTest.java b/tests/robotests/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetectorTest.java
index 386e162..19ab7d6 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetectorTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/anomaly/checker/WakeLockAnomalyDetectorTest.java
@@ -31,7 +31,6 @@
 import android.content.pm.PackageManager;
 import android.os.BatteryStats;
 import android.text.format.DateUtils;
-import android.util.ArrayMap;
 
 import com.android.internal.os.BatterySipper;
 import com.android.internal.os.BatteryStatsHelper;
@@ -64,6 +63,7 @@
     private static final int ANOMALY_UID = 111;
     private static final int NORMAL_UID = 222;
     private static final int TARGET_UID = 333;
+    private static final int INACTIVE_UID = 444;
     @Mock
     private BatteryStatsHelper mBatteryStatsHelper;
     @Mock
@@ -71,19 +71,9 @@
     @Mock
     private BatterySipper mTargetSipper;
     @Mock
-    private BatteryStats.Timer mAnomalyTimer;
-    @Mock
-    private BatteryStats.Uid.Wakelock mAnomalyWakelock;
-    @Mock
     private BatterySipper mNormalSipper;
     @Mock
-    private BatteryStats.Timer mNormalTimer;
-    @Mock
-    private BatteryStats.Timer mTargetTimer;
-    @Mock
-    private BatteryStats.Uid.Wakelock mNormalWakelock;
-    @Mock
-    private BatteryStats.Uid.Wakelock mTargetWakelock;
+    private BatterySipper mInactiveSipper;
     @Mock
     private BatteryStats.Uid mAnomalyUid;
     @Mock
@@ -91,6 +81,8 @@
     @Mock
     private BatteryStats.Uid mTargetUid;
     @Mock
+    private BatteryStats.Uid mInactiveUid;
+    @Mock
     private BatteryUtils mBatteryUtils;
     @Mock
     private PackageManager mPackageManager;
@@ -101,9 +93,6 @@
     @Mock
     private AnomalyAction mAnomalyAction;
 
-    private ArrayMap<String, BatteryStats.Uid.Wakelock> mAnomalyWakelocks;
-    private ArrayMap<String, BatteryStats.Uid.Wakelock> mNormalWakelocks;
-    private ArrayMap<String, BatteryStats.Uid.Wakelock> mTargetWakelocks;
     private WakeLockAnomalyDetector mWakelockAnomalyDetector;
     private Context mContext;
     private List<BatterySipper> mUsageList;
@@ -121,42 +110,44 @@
                 .getApplicationInfo(nullable(String.class), anyInt());
         doReturn(true).when(mAnomalyAction).isActionActive(any());
 
+        mWakelockAnomalyDetector = spy(new WakeLockAnomalyDetector(mContext, mPolicy));
+        mWakelockAnomalyDetector.mBatteryUtils = mBatteryUtils;
+        mWakelockAnomalyDetector.mAnomalyAction = mAnomalyAction;
+
         mAnomalySipper.uidObj = mAnomalyUid;
-        mAnomalyWakelocks = new ArrayMap<>();
-        mAnomalyWakelocks.put("", mAnomalyWakelock);
-        doReturn(mAnomalyWakelocks).when(mAnomalyUid).getWakelockStats();
-        doReturn(mAnomalyTimer).when(mAnomalyWakelock).getWakeTime(BatteryStats.WAKE_TYPE_PARTIAL);
+        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector)
+                .getBackgroundTotalDurationMs(eq(mAnomalyUid), anyLong());
+        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector).getCurrentDurationMs(
+                eq(mAnomalyUid), anyLong());
         doReturn(ANOMALY_UID).when(mAnomalyUid).getUid();
 
         mNormalSipper.uidObj = mNormalUid;
-        mNormalWakelocks = new ArrayMap<>();
-        mNormalWakelocks.put("", mNormalWakelock);
-        doReturn(mNormalTimer).when(mNormalWakelock).getWakeTime(BatteryStats.WAKE_TYPE_PARTIAL);
-        doReturn(mNormalWakelocks).when(mNormalUid).getWakelockStats();
+        doReturn(NORMAL_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector)
+                .getBackgroundTotalDurationMs(eq(mNormalUid), anyLong());
+        doReturn(0L).when(mWakelockAnomalyDetector).getCurrentDurationMs(eq(mNormalUid),
+                anyLong());
         doReturn(NORMAL_UID).when(mNormalUid).getUid();
 
         mTargetSipper.uidObj = mTargetUid;
-        mTargetWakelocks = new ArrayMap<>();
-        mTargetWakelocks.put("", mTargetWakelock);
-        doReturn(mTargetTimer).when(mTargetWakelock).getWakeTime(BatteryStats.WAKE_TYPE_PARTIAL);
-        doReturn(mTargetWakelocks).when(mTargetUid).getWakelockStats();
+        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector)
+                .getBackgroundTotalDurationMs(eq(mTargetUid), anyLong());
+        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector).getCurrentDurationMs(
+                eq(mTargetUid), anyLong());
         doReturn(TARGET_UID).when(mTargetUid).getUid();
 
+        mInactiveSipper.uidObj = mInactiveUid;
+        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector)
+                .getBackgroundTotalDurationMs(eq(mInactiveUid), anyLong());
+        doReturn(0L).when(mWakelockAnomalyDetector).getCurrentDurationMs(eq(mInactiveUid),
+                anyLong());
+        doReturn(INACTIVE_UID).when(mInactiveUid).getUid();
+
         mUsageList = new ArrayList<>();
         mUsageList.add(mAnomalySipper);
         mUsageList.add(mNormalSipper);
         mUsageList.add(mTargetSipper);
+        mUsageList.add(mInactiveSipper);
         doReturn(mUsageList).when(mBatteryStatsHelper).getUsageList();
-
-        mWakelockAnomalyDetector = spy(new WakeLockAnomalyDetector(mContext, mPolicy));
-        mWakelockAnomalyDetector.mBatteryUtils = mBatteryUtils;
-        mWakelockAnomalyDetector.mAnomalyAction = mAnomalyAction;
-        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector).getTotalDurationMs(
-                eq(mAnomalyTimer), anyLong());
-        doReturn(ANOMALY_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector).getTotalDurationMs(
-                eq(mTargetTimer), anyLong());
-        doReturn(NORMAL_WAKELOCK_TIME_MS).when(mWakelockAnomalyDetector).getTotalDurationMs(
-                eq(mNormalTimer), anyLong());
     }
 
     @Test