diff --git a/src/com/android/settings/fuelgauge/BatteryInfo.java b/src/com/android/settings/fuelgauge/BatteryInfo.java
index ce441c8..af7dded 100644
--- a/src/com/android/settings/fuelgauge/BatteryInfo.java
+++ b/src/com/android/settings/fuelgauge/BatteryInfo.java
@@ -375,8 +375,8 @@
         boolean first = true;
         final BatteryStatsHistoryIterator iterator1 =
                 mBatteryUsageStats.iterateBatteryStatsHistory();
-        final HistoryItem rec = new HistoryItem();
-        while (iterator1.next(rec)) {
+        HistoryItem rec;
+        while ((rec = iterator1.next()) != null) {
             pos++;
             if (first) {
                 first = false;
@@ -420,7 +420,7 @@
         if (endWalltime > startWalltime) {
             final BatteryStatsHistoryIterator iterator2 =
                     mBatteryUsageStats.iterateBatteryStatsHistory();
-            while (iterator2.next(rec) && i < N) {
+            while ((rec = iterator2.next()) != null && i < N) {
                 if (rec.isDeltaData()) {
                     curWalltime += rec.time - lastRealtime;
                     lastRealtime = rec.time;
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java b/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
index 284dbe3..e7a9811 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
@@ -53,8 +53,6 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.mockito.invocation.InvocationOnMock;
-import org.mockito.stubbing.Answer;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
@@ -333,28 +331,23 @@
         // Mock out new data every time iterateBatteryStatsHistory is called.
         doAnswer(invocation -> {
             BatteryStatsHistoryIterator iterator = mock(BatteryStatsHistoryIterator.class);
-            doAnswer(new Answer<Boolean>() {
-                private int mCount = 0;
-                private final long[] mTimes = {1000, 1500, 2000};
-                private final byte[] mLevels = {99, 98, 97};
-
-                @Override
-                public Boolean answer(InvocationOnMock invocation) throws Throwable {
-                    if (mCount == mTimes.length) {
-                        return false;
-                    }
-                    BatteryStats.HistoryItem record = invocation.getArgument(0);
-                    record.cmd = BatteryStats.HistoryItem.CMD_UPDATE;
-                    record.time = mTimes[mCount];
-                    record.batteryLevel = mLevels[mCount];
-                    mCount++;
-                    return true;
-                }
-            }).when(iterator).next(any(BatteryStats.HistoryItem.class));
+            when(iterator.next()).thenReturn(
+                    makeHistoryIterm(1000, 99),
+                    makeHistoryIterm(1500, 98),
+                    makeHistoryIterm(2000, 97),
+                    null);
             return iterator;
         }).when(mBatteryUsageStats).iterateBatteryStatsHistory();
     }
 
+    private BatteryStats.HistoryItem makeHistoryIterm(long time, int batteryLevel) {
+        BatteryStats.HistoryItem record = new BatteryStats.HistoryItem();
+        record.cmd = BatteryStats.HistoryItem.CMD_UPDATE;
+        record.time = time;
+        record.batteryLevel = (byte) batteryLevel;
+        return record;
+    }
+
     private void assertOnlyHistory(BatteryInfo info) {
         mockBatteryStatsHistory();
         UsageView view = mock(UsageView.class);