Copy Battery Usage Database from SettingsIntelligence to Settings.

Bug: 253395332
Test: make RunSettingsRoboTests + manually
Change-Id: Ibdd2ace10d9e0893b3d96b345d563307b1890df6
diff --git a/Android.bp b/Android.bp
index 9ddadaf..f980bfa 100644
--- a/Android.bp
+++ b/Android.bp
@@ -90,8 +90,11 @@
         "WifiTrackerLib",
         "SettingsLibActivityEmbedding",
         "Settings-change-ids",
+        "androidx.room_room-runtime",
     ],
 
+    plugins: ["androidx.room_room-compiler-plugin"],
+
     libs: [
         "telephony-common",
         "ims-common",
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryState.java b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryState.java
new file mode 100644
index 0000000..11db118
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryState.java
@@ -0,0 +1,397 @@
+/*
+ * 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.db;
+
+import android.content.ContentValues;
+import android.content.Intent;
+import android.os.BatteryManager;
+
+import androidx.room.Entity;
+import androidx.room.PrimaryKey;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+
+import java.text.SimpleDateFormat;
+import java.time.Duration;
+import java.util.Date;
+import java.util.Locale;
+import java.util.TimeZone;
+
+/** A {@link Entity} class to save battery states snapshot into database. */
+@Entity
+public class BatteryState {
+    private static String sCacheZoneId;
+    private static SimpleDateFormat sCacheSimpleDateFormat;
+
+    @PrimaryKey(autoGenerate = true)
+    private long mId;
+
+    // Records the app relative information.
+    public final long uid;
+    public final long userId;
+    public final String appLabel;
+    public final String packageName;
+    // Whether the data is represented as system component or not?
+    public final boolean isHidden;
+    // Records the timestamp relative information.
+    public final long bootTimestamp;
+    public final long timestamp;
+    public final String zoneId;
+    // Records the battery usage relative information.
+    public final double totalPower;
+    public final double consumePower;
+    public final double percentOfTotal;
+    public final long foregroundUsageTimeInMs;
+    public final long backgroundUsageTimeInMs;
+    public final int drainType;
+    public final int consumerType;
+    // Records the battery intent relative information.
+    public final int batteryLevel;
+    public final int batteryStatus;
+    public final int batteryHealth;
+
+    public BatteryState(
+            long uid,
+            long userId,
+            String appLabel,
+            String packageName,
+            boolean isHidden,
+            long bootTimestamp,
+            long timestamp,
+            String zoneId,
+            double totalPower,
+            double consumePower,
+            double percentOfTotal,
+            long foregroundUsageTimeInMs,
+            long backgroundUsageTimeInMs,
+            int drainType,
+            int consumerType,
+            int batteryLevel,
+            int batteryStatus,
+            int batteryHealth) {
+        // Records the app relative information.
+        this.uid = uid;
+        this.userId = userId;
+        this.appLabel = appLabel;
+        this.packageName = packageName;
+        this.isHidden = isHidden;
+        // Records the timestamp relative information.
+        this.bootTimestamp = bootTimestamp;
+        this.timestamp = timestamp;
+        this.zoneId = zoneId;
+        // Records the battery usage relative information.
+        this.totalPower = totalPower;
+        this.consumePower = consumePower;
+        this.percentOfTotal = percentOfTotal;
+        this.foregroundUsageTimeInMs = foregroundUsageTimeInMs;
+        this.backgroundUsageTimeInMs = backgroundUsageTimeInMs;
+        this.drainType = drainType;
+        this.consumerType = consumerType;
+        // Records the battery intent relative information.
+        this.batteryLevel = batteryLevel;
+        this.batteryStatus = batteryStatus;
+        this.batteryHealth = batteryHealth;
+    }
+
+    /** Sets the auto-generated content ID. */
+    public void setId(long id) {
+        this.mId = id;
+    }
+
+    /** Gets the auto-generated content ID. */
+    public long getId() {
+        return mId;
+    }
+
+    @Override
+    @SuppressWarnings("JavaUtilDate")
+    public String toString() {
+        final String currentZoneId = TimeZone.getDefault().getID();
+        if (!currentZoneId.equals(sCacheZoneId) || sCacheSimpleDateFormat == null) {
+            sCacheZoneId = currentZoneId;
+            sCacheSimpleDateFormat = new SimpleDateFormat("MMM dd,yyyy HH:mm:ss", Locale.US);
+        }
+        final String recordAtDateTime = sCacheSimpleDateFormat.format(new Date(timestamp));
+        final StringBuilder builder = new StringBuilder()
+                .append("\nBatteryState{")
+                .append(String.format(Locale.US,
+                        "\n\tpackage=%s|label=%s|uid=%d|userId=%d|isHidden=%b",
+                        packageName, appLabel, uid, userId, isHidden))
+                .append(String.format(Locale.US, "\n\ttimestamp=%s|zoneId=%s|bootTimestamp=%d",
+                        recordAtDateTime, zoneId, Duration.ofMillis(bootTimestamp).getSeconds()))
+                .append(String.format(Locale.US,
+                        "\n\tusage=%f|total=%f|consume=%f|elapsedTime=%d|%d",
+                        percentOfTotal, totalPower, consumePower,
+                        Duration.ofMillis(foregroundUsageTimeInMs).getSeconds(),
+                        Duration.ofMillis(backgroundUsageTimeInMs).getSeconds()))
+                .append(String.format(Locale.US,
+                        "\n\tdrain=%d|consumer=%d", drainType, consumerType))
+                .append(String.format(Locale.US, "\n\tbattery=%d|status=%d|health=%d\n}",
+                        batteryLevel, batteryStatus, batteryHealth));
+        return builder.toString();
+    }
+
+
+    /** Creates new {@link BatteryState} from {@link ContentValues}. */
+    public static BatteryState create(ContentValues contentValues) {
+        Builder builder = BatteryState.newBuilder();
+        if (contentValues.containsKey("uid")) {
+            builder.setUid(contentValues.getAsLong("uid"));
+        }
+        if (contentValues.containsKey("userId")) {
+            builder.setUserId(contentValues.getAsLong("userId"));
+        }
+        if (contentValues.containsKey("appLabel")) {
+            builder.setAppLabel(contentValues.getAsString("appLabel"));
+        }
+        if (contentValues.containsKey("packageName")) {
+            builder.setPackageName(contentValues.getAsString("packageName"));
+        }
+        if (contentValues.containsKey("isHidden")) {
+            builder.setIsHidden(contentValues.getAsBoolean("isHidden"));
+        }
+        if (contentValues.containsKey("bootTimestamp")) {
+            builder.setBootTimestamp(contentValues.getAsLong("bootTimestamp"));
+        }
+        if (contentValues.containsKey("timestamp")) {
+            builder.setTimestamp(contentValues.getAsLong("timestamp"));
+        }
+        if (contentValues.containsKey("consumePower")) {
+            builder.setConsumePower(contentValues.getAsDouble("consumePower"));
+        }
+        if (contentValues.containsKey("totalPower")) {
+            builder.setTotalPower(contentValues.getAsDouble("totalPower"));
+        }
+        if (contentValues.containsKey("percentOfTotal")) {
+            builder.setPercentOfTotal(contentValues.getAsDouble("percentOfTotal"));
+        }
+        if (contentValues.containsKey("foregroundUsageTimeInMs")) {
+            builder.setForegroundUsageTimeInMs(
+                    contentValues.getAsLong("foregroundUsageTimeInMs"));
+        }
+        if (contentValues.containsKey("backgroundUsageTimeInMs")) {
+            builder.setBackgroundUsageTimeInMs(
+                    contentValues.getAsLong("backgroundUsageTimeInMs"));
+        }
+        if (contentValues.containsKey("drainType")) {
+            builder.setDrainType(contentValues.getAsInteger("drainType"));
+        }
+        if (contentValues.containsKey("consumerType")) {
+            builder.setConsumerType(contentValues.getAsInteger("consumerType"));
+        }
+        if (contentValues.containsKey("batteryLevel")) {
+            builder.setBatteryLevel(contentValues.getAsInteger("batteryLevel"));
+        }
+        if (contentValues.containsKey("batteryStatus")) {
+            builder.setBatteryStatus(contentValues.getAsInteger("batteryStatus"));
+        }
+        if (contentValues.containsKey("batteryHealth")) {
+            builder.setBatteryHealth(contentValues.getAsInteger("batteryHealth"));
+        }
+        return builder.build();
+    }
+
+    /** Creates a new {@link Builder} instance. */
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    /** A convenience builder class to improve readability. */
+    public static class Builder {
+        private long mUid;
+        private long mUserId;
+        private String mAppLabel;
+        private String mPackageName;
+        private boolean mIsHidden;
+        private long mBootTimestamp;
+        private long mTimestamp;
+        private double mTotalPower;
+        private double mConsumePower;
+        private double mPercentOfTotal;
+        private long mForegroundUsageTimeInMs;
+        private long mBackgroundUsageTimeInMs;
+        private int mDrainType;
+        private int mConsumerType;
+        private int mBatteryLevel;
+        private int mBatteryStatus;
+        private int mBatteryHealth;
+
+        /** Sets the uid. */
+        @CanIgnoreReturnValue
+        public Builder setUid(long uid) {
+            this.mUid = uid;
+            return this;
+        }
+
+        /** Sets the user ID. */
+        @CanIgnoreReturnValue
+        public Builder setUserId(long userId) {
+            this.mUserId = userId;
+            return this;
+        }
+
+        /** Sets the app label. */
+        @CanIgnoreReturnValue
+        public Builder setAppLabel(String appLabel) {
+            this.mAppLabel = appLabel;
+            return this;
+        }
+
+        /** Sets the package name. */
+        @CanIgnoreReturnValue
+        public Builder setPackageName(String packageName) {
+            this.mPackageName = packageName;
+            return this;
+        }
+
+        /** Sets the is hidden value. */
+        @CanIgnoreReturnValue
+        public Builder setIsHidden(boolean isHidden) {
+            this.mIsHidden = isHidden;
+            return this;
+        }
+
+        /** Sets the boot timestamp. */
+        @CanIgnoreReturnValue
+        public Builder setBootTimestamp(long bootTimestamp) {
+            this.mBootTimestamp = bootTimestamp;
+            return this;
+        }
+
+        /** Sets the timestamp. */
+        @CanIgnoreReturnValue
+        public Builder setTimestamp(long timestamp) {
+            this.mTimestamp = timestamp;
+            return this;
+        }
+
+        /** Sets the total power. */
+        @CanIgnoreReturnValue
+        public Builder setTotalPower(double totalPower) {
+            this.mTotalPower = totalPower;
+            return this;
+        }
+
+        /** Sets the consumed power. */
+        @CanIgnoreReturnValue
+        public Builder setConsumePower(double consumePower) {
+            this.mConsumePower = consumePower;
+            return this;
+        }
+
+        /** Sets the percentage of total. */
+        @CanIgnoreReturnValue
+        public Builder setPercentOfTotal(double percentOfTotal) {
+            this.mPercentOfTotal = percentOfTotal;
+            return this;
+        }
+
+        /** Sets the foreground usage time. */
+        @CanIgnoreReturnValue
+        public Builder setForegroundUsageTimeInMs(long foregroundUsageTimeInMs) {
+            this.mForegroundUsageTimeInMs = foregroundUsageTimeInMs;
+            return this;
+        }
+
+        /** Sets the background usage time. */
+        @CanIgnoreReturnValue
+        public Builder setBackgroundUsageTimeInMs(long backgroundUsageTimeInMs) {
+            this.mBackgroundUsageTimeInMs = backgroundUsageTimeInMs;
+            return this;
+        }
+
+        /** Sets the drain type. */
+        @CanIgnoreReturnValue
+        public Builder setDrainType(int drainType) {
+            this.mDrainType = drainType;
+            return this;
+        }
+
+        /** Sets the consumer type. */
+        @CanIgnoreReturnValue
+        public Builder setConsumerType(int consumerType) {
+            this.mConsumerType = consumerType;
+            return this;
+        }
+
+        /** Sets the battery level. */
+        @CanIgnoreReturnValue
+        public Builder setBatteryLevel(int batteryLevel) {
+            this.mBatteryLevel = batteryLevel;
+            return this;
+        }
+
+        /** Sets the battery status. */
+        @CanIgnoreReturnValue
+        public Builder setBatteryStatus(int batteryStatus) {
+            this.mBatteryStatus = batteryStatus;
+            return this;
+        }
+
+        /** Sets the battery health. */
+        @CanIgnoreReturnValue
+        public Builder setBatteryHealth(int batteryHealth) {
+            this.mBatteryHealth = batteryHealth;
+            return this;
+        }
+
+        /** Sets the battery intent. */
+        @CanIgnoreReturnValue
+        public Builder setBatteryIntent(Intent batteryIntent) {
+            final int level = batteryIntent.getIntExtra(BatteryManager.EXTRA_LEVEL, -1);
+            final int scale = batteryIntent.getIntExtra(BatteryManager.EXTRA_SCALE, 0);
+            this.mBatteryLevel =
+                    scale == 0
+                            ? -1 /*invalid battery level*/
+                            : Math.round((level / (float) scale) * 100f);
+            this.mBatteryStatus =
+                    batteryIntent.getIntExtra(
+                            BatteryManager.EXTRA_STATUS,
+                            BatteryManager.BATTERY_STATUS_UNKNOWN);
+            this.mBatteryHealth =
+                    batteryIntent.getIntExtra(
+                            BatteryManager.EXTRA_HEALTH,
+                            BatteryManager.BATTERY_HEALTH_UNKNOWN);
+            return this;
+        }
+
+        /** Builds the BatteryState. */
+        public BatteryState build() {
+            return new BatteryState(
+                    mUid,
+                    mUserId,
+                    mAppLabel,
+                    mPackageName,
+                    mIsHidden,
+                    mBootTimestamp,
+                    mTimestamp,
+                    /*zoneId=*/ TimeZone.getDefault().getID(),
+                    mTotalPower,
+                    mConsumePower,
+                    mPercentOfTotal,
+                    mForegroundUsageTimeInMs,
+                    mBackgroundUsageTimeInMs,
+                    mDrainType,
+                    mConsumerType,
+                    mBatteryLevel,
+                    mBatteryStatus,
+                    mBatteryHealth);
+        }
+
+        private Builder() {}
+    }
+}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDao.java b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDao.java
new file mode 100644
index 0000000..b1afa6b
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDao.java
@@ -0,0 +1,63 @@
+/*
+ * 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.db;
+
+import android.database.Cursor;
+
+import androidx.room.Dao;
+import androidx.room.Insert;
+import androidx.room.OnConflictStrategy;
+import androidx.room.Query;
+
+import java.util.List;
+
+/** Data access object for accessing {@link BatteryState} in the database. */
+@Dao
+public interface BatteryStateDao {
+
+    /** Inserts a {@link BatteryState} data into the database. */
+    @Insert(onConflict = OnConflictStrategy.REPLACE)
+    void insert(BatteryState state);
+
+    /** Inserts {@link BatteryState} data into the database. */
+    @Insert(onConflict = OnConflictStrategy.REPLACE)
+    void insertAll(List<BatteryState> states);
+
+    /** Lists all recorded data after a specific timestamp. */
+    @Query("SELECT * FROM BatteryState WHERE timestamp > :timestamp ORDER BY timestamp DESC")
+    List<BatteryState> getAllAfter(long timestamp);
+
+    /** Gets the {@link Cursor} of all recorded data from a specific timestamp. */
+    @Query("SELECT * FROM BatteryState WHERE timestamp >= :timestamp ORDER BY timestamp DESC")
+    Cursor getCursorAfter(long timestamp);
+
+    /** Get the count of distinct timestamp after a specific timestamp. */
+    @Query("SELECT COUNT(DISTINCT timestamp) FROM BatteryState WHERE timestamp > :timestamp")
+    int getDistinctTimestampCount(long timestamp);
+
+    /** Lists all distinct timestamps after a specific timestamp. */
+    @Query("SELECT DISTINCT timestamp FROM BatteryState WHERE timestamp > :timestamp")
+    List<Long> getDistinctTimestamps(long timestamp);
+
+    /** Deletes all recorded data before a specific timestamp. */
+    @Query("DELETE FROM BatteryState WHERE timestamp <= :timestamp")
+    void clearAllBefore(long timestamp);
+
+    /** Clears all recorded data in the database. */
+    @Query("DELETE FROM BatteryState")
+    void clearAll();
+}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java
new file mode 100644
index 0000000..9396546
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java
@@ -0,0 +1,57 @@
+/*
+ * 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.db;
+
+import android.content.Context;
+import android.util.Log;
+
+import androidx.room.Database;
+import androidx.room.Room;
+import androidx.room.RoomDatabase;
+
+/** A {@link RoomDatabase} for battery usage states history. */
+@Database(
+        entities = {BatteryState.class},
+        version = 1)
+public abstract class BatteryStateDatabase extends RoomDatabase {
+    private static final String TAG = "BatteryStateDatabase";
+
+    private static BatteryStateDatabase sBatteryStateDatabase;
+
+    /** Provides DAO for battery state table. */
+    public abstract BatteryStateDao batteryStateDao();
+
+    /** Gets or creates an instance of {@link RoomDatabase}. */
+    public static BatteryStateDatabase getInstance(Context context) {
+        if (sBatteryStateDatabase == null) {
+            sBatteryStateDatabase =
+                    Room.databaseBuilder(
+                                    context, BatteryStateDatabase.class, "battery-usage-db-v1")
+                            // Allows accessing data in the main thread for dumping bugreport.
+                            .allowMainThreadQueries()
+                            .fallbackToDestructiveMigration()
+                            .build();
+            Log.d(TAG, "initialize battery states database");
+        }
+        return sBatteryStateDatabase;
+    }
+
+    /** Sets the instance of {@link RoomDatabase}. */
+    public static void setBatteryStateDatabase(BatteryStateDatabase database) {
+        BatteryStateDatabase.sBatteryStateDatabase = database;
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDaoTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDaoTest.java
new file mode 100644
index 0000000..41e3f4d
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDaoTest.java
@@ -0,0 +1,135 @@
+/*
+ * 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.db;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.database.Cursor;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settings.testutils.BatteryTestUtils;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+
+import java.util.List;
+
+/** Tests for {@link BatteryStateDao}. */
+@RunWith(RobolectricTestRunner.class)
+public final class BatteryStateDaoTest {
+    private static final int CURSOR_COLUMN_SIZE = 19;
+    private static final long TIMESTAMP1 = System.currentTimeMillis();
+    private static final long TIMESTAMP2 = System.currentTimeMillis() + 2;
+    private static final long TIMESTAMP3 = System.currentTimeMillis() + 4;
+    private static final String PACKAGE_NAME1 = "com.android.apps.settings";
+    private static final String PACKAGE_NAME2 = "com.android.apps.calendar";
+    private static final String PACKAGE_NAME3 = "com.android.apps.gmail";
+
+    private Context mContext;
+    private BatteryStateDatabase mDatabase;
+    private BatteryStateDao mBatteryStateDao;
+
+    @Before
+    public void setUp() {
+        mContext = ApplicationProvider.getApplicationContext();
+        mDatabase = BatteryTestUtils.setUpBatteryStateDatabase(mContext);
+        mBatteryStateDao = mDatabase.batteryStateDao();
+        BatteryTestUtils.insertDataToBatteryStateDatabase(mContext, TIMESTAMP3, PACKAGE_NAME3);
+        BatteryTestUtils.insertDataToBatteryStateDatabase(mContext, TIMESTAMP2, PACKAGE_NAME2);
+        BatteryTestUtils.insertDataToBatteryStateDatabase(
+                mContext, TIMESTAMP1, PACKAGE_NAME1, /*multiple=*/ true);
+    }
+
+    @After
+    public void closeDb() {
+        mDatabase.close();
+        BatteryStateDatabase.setBatteryStateDatabase(/*database=*/ null);
+    }
+
+    @Test
+    public void batteryStateDao_insertAll() throws Exception {
+        final List<BatteryState> states = mBatteryStateDao.getAllAfter(TIMESTAMP1);
+        assertThat(states).hasSize(2);
+        // Verifies the queried battery states.
+        assertBatteryState(states.get(0), TIMESTAMP3, PACKAGE_NAME3);
+        assertBatteryState(states.get(1), TIMESTAMP2, PACKAGE_NAME2);
+    }
+
+    @Test
+    public void batteryStateDao_getCursorAfter() throws Exception {
+        final Cursor cursor = mBatteryStateDao.getCursorAfter(TIMESTAMP2);
+        assertThat(cursor.getCount()).isEqualTo(2);
+        assertThat(cursor.getColumnCount()).isEqualTo(CURSOR_COLUMN_SIZE);
+        // Verifies the queried first battery state.
+        cursor.moveToFirst();
+        assertThat(cursor.getString(4 /*packageName*/)).isEqualTo(PACKAGE_NAME3);
+        // Verifies the queried second battery state.
+        cursor.moveToNext();
+        assertThat(cursor.getString(4 /*packageName*/)).isEqualTo(PACKAGE_NAME2);
+    }
+
+    @Test
+    public void batteryStateDao_clearAllBefore() throws Exception {
+        mBatteryStateDao.clearAllBefore(TIMESTAMP2);
+
+        final List<BatteryState> states = mBatteryStateDao.getAllAfter(0);
+        assertThat(states).hasSize(1);
+        // Verifies the queried battery state.
+        assertBatteryState(states.get(0), TIMESTAMP3, PACKAGE_NAME3);
+    }
+
+    @Test
+    public void batteryStateDao_clearAll() throws Exception {
+        assertThat(mBatteryStateDao.getAllAfter(0)).hasSize(3);
+        mBatteryStateDao.clearAll();
+        assertThat(mBatteryStateDao.getAllAfter(0)).isEmpty();
+    }
+
+    @Test
+    public void getInstance_createNewInstance() throws Exception {
+        BatteryStateDatabase.setBatteryStateDatabase(/*database=*/ null);
+        assertThat(BatteryStateDatabase.getInstance(mContext)).isNotNull();
+    }
+
+    @Test
+    public void getDistinctTimestampCount_returnsExpectedResult() {
+        assertThat(mBatteryStateDao.getDistinctTimestampCount(/*timestamp=*/ 0))
+                .isEqualTo(3);
+        assertThat(mBatteryStateDao.getDistinctTimestampCount(TIMESTAMP1))
+                .isEqualTo(2);
+    }
+
+    @Test
+    public void getDistinctTimestamps_returnsExpectedResult() {
+        final List<Long> timestamps =
+                mBatteryStateDao.getDistinctTimestamps(/*timestamp=*/ 0);
+
+        assertThat(timestamps).hasSize(3);
+        assertThat(timestamps).containsExactly(TIMESTAMP1, TIMESTAMP2, TIMESTAMP3);
+    }
+
+    private static void assertBatteryState(
+            BatteryState state, long timestamp, String packageName) {
+        assertThat(state.timestamp).isEqualTo(timestamp);
+        assertThat(state.packageName).isEqualTo(packageName);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateTest.java
new file mode 100644
index 0000000..ef23c41
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateTest.java
@@ -0,0 +1,98 @@
+/*
+ * 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.db;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Intent;
+import android.os.BatteryManager;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+
+/** Tests for {@link BatteryState}. */
+@RunWith(RobolectricTestRunner.class)
+public final class BatteryStateTest {
+    private static final int BATTERY_LEVEL = 45;
+    private static final int BATTERY_STATUS = BatteryManager.BATTERY_STATUS_FULL;
+    private static final int BATTERY_HEALTH = BatteryManager.BATTERY_HEALTH_COLD;
+
+    private Intent mBatteryIntent;
+
+    @Before
+    public void setUp() {
+        mBatteryIntent = new Intent(Intent.ACTION_BATTERY_CHANGED);
+        // Inserts the battery states into intent.
+        mBatteryIntent.putExtra(BatteryManager.EXTRA_LEVEL, BATTERY_LEVEL);
+        mBatteryIntent.putExtra(BatteryManager.EXTRA_STATUS, BATTERY_STATUS);
+        mBatteryIntent.putExtra(BatteryManager.EXTRA_HEALTH, BATTERY_HEALTH);
+    }
+
+    @Test
+    public void testBuilder_returnsExpectedResult() {
+        mBatteryIntent.putExtra(BatteryManager.EXTRA_SCALE, 100);
+        BatteryState state = create(mBatteryIntent);
+
+        // Verifies the app relative information.
+        assertThat(state.uid).isEqualTo(1001L);
+        assertThat(state.userId).isEqualTo(100L);
+        assertThat(state.appLabel).isEqualTo("Settings");
+        assertThat(state.packageName).isEqualTo("com.android.settings");
+        assertThat(state.isHidden).isTrue();
+        assertThat(state.bootTimestamp).isEqualTo(101L);
+        assertThat(state.timestamp).isEqualTo(100001L);
+        // Verifies the battery relative information.
+        assertThat(state.totalPower).isEqualTo(100);
+        assertThat(state.consumePower).isEqualTo(3);
+        assertThat(state.percentOfTotal).isEqualTo(10);
+        assertThat(state.foregroundUsageTimeInMs).isEqualTo(60000);
+        assertThat(state.backgroundUsageTimeInMs).isEqualTo(10000);
+        assertThat(state.drainType).isEqualTo(1);
+        assertThat(state.consumerType).isEqualTo(2);
+        assertThat(state.batteryLevel).isEqualTo(BATTERY_LEVEL);
+        assertThat(state.batteryStatus).isEqualTo(BATTERY_STATUS);
+        assertThat(state.batteryHealth).isEqualTo(BATTERY_HEALTH);
+    }
+
+    @Test
+    public void create_withoutBatteryScale_returnsStateWithInvalidLevel() {
+        BatteryState state = create(mBatteryIntent);
+        assertThat(state.batteryLevel).isEqualTo(-1);
+    }
+
+    private static BatteryState create(Intent intent) {
+        return BatteryState.newBuilder()
+                .setUid(1001L)
+                .setUserId(100L)
+                .setAppLabel("Settings")
+                .setPackageName("com.android.settings")
+                .setIsHidden(true)
+                .setBootTimestamp(101L)
+                .setTimestamp(100001L)
+                .setTotalPower(100f)
+                .setConsumePower(3f)
+                .setPercentOfTotal(10f)
+                .setForegroundUsageTimeInMs(60000)
+                .setBackgroundUsageTimeInMs(10000)
+                .setDrainType(1)
+                .setConsumerType(2)
+                .setBatteryIntent(intent)
+                .build();
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/testutils/BatteryTestUtils.java b/tests/robotests/src/com/android/settings/testutils/BatteryTestUtils.java
index e4e26d2..fa3ee10 100644
--- a/tests/robotests/src/com/android/settings/testutils/BatteryTestUtils.java
+++ b/tests/robotests/src/com/android/settings/testutils/BatteryTestUtils.java
@@ -16,8 +16,21 @@
 
 package com.android.settings.testutils;
 
+import android.content.Context;
 import android.content.Intent;
 import android.os.BatteryManager;
+import android.os.UserManager;
+
+import androidx.room.Room;
+
+import com.android.settings.fuelgauge.batteryusage.db.BatteryState;
+import com.android.settings.fuelgauge.batteryusage.db.BatteryStateDao;
+import com.android.settings.fuelgauge.batteryusage.db.BatteryStateDatabase;
+
+import com.google.common.collect.ImmutableList;
+
+import org.robolectric.Shadows;
+import org.robolectric.shadows.ShadowUserManager;
 
 public class BatteryTestUtils {
 
@@ -37,6 +50,65 @@
                 BatteryManager.BATTERY_STATUS_DISCHARGING);
     }
 
+    /** Sets the work profile mode. */
+    public static void setWorkProfile(Context context) {
+        final UserManager userManager = context.getSystemService(UserManager.class);
+        Shadows.shadowOf(userManager).setManagedProfile(true);
+        // Changes out of the default system user so isSystemUser() returns false.
+        final int userId = 1001;
+        Shadows.shadowOf(userManager)
+                .addUser(userId, "name", /*flags=*/ ShadowUserManager.FLAG_PRIMARY);
+        Shadows.shadowOf(userManager).switchUser(userId);
+    }
+
+    /** Creates and sets up the in-memory {@link BatteryStateDatabase}. */
+    public static BatteryStateDatabase setUpBatteryStateDatabase(Context context) {
+        final BatteryStateDatabase inMemoryDatabase =
+                Room.inMemoryDatabaseBuilder(context, BatteryStateDatabase.class)
+                        .allowMainThreadQueries()
+                        .build();
+        BatteryStateDatabase.setBatteryStateDatabase(inMemoryDatabase);
+        return inMemoryDatabase;
+    }
+
+    /** Inserts a fake data into the database for testing. */
+    public static void insertDataToBatteryStateDatabase(
+            Context context, long timestamp, String packageName) {
+        insertDataToBatteryStateDatabase(context, timestamp, packageName, /*multiple=*/ false);
+    }
+
+    /** Inserts a fake data into the database for testing. */
+    public static void insertDataToBatteryStateDatabase(
+            Context context, long timestamp, String packageName, boolean multiple) {
+        final BatteryState state =
+                new BatteryState(
+                        /*uid=*/ 1001L,
+                        /*userId=*/ 100L,
+                        /*appLabel=*/ "Settings",
+                        packageName,
+                        /*isHidden=*/ true,
+                        /*bootTimestamp=*/ timestamp - 1,
+                        timestamp,
+                        /*zoneId=*/ "Europe/Paris",
+                        /*totalPower=*/ 100f,
+                        /*consumePower=*/ 0.3f,
+                        /*percentOfTotal=*/ 10f,
+                        /*foregroundUsageTimeInMs=*/ 60000,
+                        /*backgroundUsageTimeInMs=*/ 10000,
+                        /*drainType=*/ 1,
+                        /*consumerType=*/ 2,
+                        /*batteryLevel=*/ 31,
+                        /*batteryStatus=*/ 0,
+                        /*batteryHealth=*/ 0);
+        BatteryStateDao dao =
+                BatteryStateDatabase.getInstance(context).batteryStateDao();
+        if (multiple) {
+            dao.insertAll(ImmutableList.of(state));
+        } else {
+            dao.insert(state);
+        }
+    }
+
     private static Intent getCustomBatteryIntent(int plugged, int level, int scale, int status) {
         Intent intent = new Intent();
         intent.putExtra(BatteryManager.EXTRA_PLUGGED, plugged);