Add new database table and proto for battery usage reattribution (2/5)

Bug: 346706894
Test: atest SettingsRoboTests:com.android.settings.fuelgauge.batteryusage
Flag: EXEMPT bug fix
Change-Id: If360246d974abdea7004023aedcf1a4be7b63633
diff --git a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
index 1c6ff54..8f4d4dd 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/ConvertUtils.java
@@ -201,6 +201,20 @@
         return defaultInstance;
     }
 
+    /** Gets the encoded string from {@link BatteryReattribute} instance. */
+    @NonNull
+    public static String encodeBatteryReattribute(
+            @NonNull BatteryReattribute batteryReattribute) {
+        return Base64.encodeToString(batteryReattribute.toByteArray(), Base64.DEFAULT);
+    }
+
+    /** Gets the decoded {@link BatteryReattribute} instance from string. */
+    @NonNull
+    public static BatteryReattribute decodeBatteryReattribute(@NonNull String content) {
+        return BatteryUtils.parseProtoFromString(
+                content, BatteryReattribute.getDefaultInstance());
+    }
+
     /** Converts to {@link BatteryHistEntry} */
     public static BatteryHistEntry convertToBatteryHistEntry(
             BatteryEntry entry, BatteryUsageStats batteryUsageStats) {
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java b/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
index 0bb6286..7620323 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DatabaseUtils.java
@@ -429,6 +429,7 @@
                         database.batteryEventDao().clearAll();
                         database.batteryStateDao().clearAll();
                         database.batteryUsageSlotDao().clearAll();
+                        database.batteryReattributeDao().clearAll();
                     } catch (RuntimeException e) {
                         Log.e(TAG, "clearAll() failed", e);
                     }
@@ -446,6 +447,7 @@
                         database.batteryEventDao().clearAllAfter(startTimestamp);
                         database.batteryStateDao().clearAllAfter(startTimestamp);
                         database.batteryUsageSlotDao().clearAllAfter(startTimestamp);
+                        database.batteryReattributeDao().clearAllAfter(startTimestamp);
                     } catch (RuntimeException e) {
                         Log.e(TAG, "clearAllAfter() failed", e);
                     }
@@ -466,6 +468,7 @@
                         database.batteryEventDao().clearAllBefore(earliestTimestamp);
                         database.batteryStateDao().clearAllBefore(earliestTimestamp);
                         database.batteryUsageSlotDao().clearAllBefore(earliestTimestamp);
+                        database.batteryReattributeDao().clearAllBefore(earliestTimestamp);
                     } catch (RuntimeException e) {
                         Log.e(TAG, "clearAllBefore() failed", e);
                     }
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDao.kt b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDao.kt
new file mode 100644
index 0000000..79c9d00
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDao.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2024 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 androidx.room.Dao;
+import androidx.room.Insert;
+import androidx.room.OnConflictStrategy;
+import androidx.room.Query;
+
+import java.util.List;
+
+/** DAO for accessing {@link BatteryReattributeEntity} in the database. */
+@Dao
+public interface BatteryReattributeDao {
+
+    /** Inserts a {@link BatteryReattributeEntity} data into the database. */
+    @Insert(onConflict = OnConflictStrategy.REPLACE)
+    fun insert(event: BatteryReattributeEntity)
+
+    /** Gets all recorded data after a specific timestamp. */
+    @Query(
+            "SELECT * FROM BatteryReattributeEntity WHERE "
+                    + "timestampStart >= :timestampStart ORDER BY timestampStart DESC")
+    fun getAllAfter(timestampStart: Long): List<BatteryReattributeEntity>
+
+    /** Deletes all recorded data before a specific timestamp. */
+    @Query("DELETE FROM BatteryReattributeEntity WHERE timestampStart <= :timestampStart")
+    fun clearAllBefore(timestampStart: Long)
+
+    /** Deletes all recorded data after a specific timestamp. */
+    @Query("DELETE FROM BatteryReattributeEntity WHERE timestampStart >= :timestampStart")
+    fun clearAllAfter(timestampStart: Long)
+
+    /** Clears all recorded data in the database. */
+    @Query("DELETE FROM BatteryReattributeEntity") fun clearAll()
+}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntity.java b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntity.java
new file mode 100644
index 0000000..aa7e50e
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntity.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 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.android.settings.fuelgauge.batteryusage.ConvertUtils.utcToLocalTimeForLogging;
+
+import com.android.settings.fuelgauge.batteryusage.BatteryReattribute;
+import com.android.settings.fuelgauge.batteryusage.ConvertUtils;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
+import androidx.room.Entity;
+import androidx.room.PrimaryKey;
+
+/** A {@link Entity} for battery usage reattribution data in the database. */
+@Entity
+public class BatteryReattributeEntity {
+
+    /** The start timestamp of this record data. */
+    @PrimaryKey
+    public final long timestampStart;
+
+    /** The end timestamp of this record data. */
+    public final long timestampEnd;
+
+    /** The battery usage reattribution data for corresponding  uids. */
+    public final String reattributeData;
+
+    public BatteryReattributeEntity(@NonNull BatteryReattribute batteryReattribute) {
+        this(
+                batteryReattribute.getTimestampStart(),
+                batteryReattribute.getTimestampEnd(),
+                ConvertUtils.encodeBatteryReattribute(batteryReattribute));
+    }
+
+    @VisibleForTesting
+    BatteryReattributeEntity(
+            long timestampStart, long timestampEnd, @NonNull String reattributeData) {
+        this.timestampStart = timestampStart;
+        this.timestampEnd = timestampEnd;
+        this.reattributeData = reattributeData;
+    }
+
+    @NonNull
+    @Override
+    public String toString() {
+        final StringBuilder builder = new StringBuilder()
+                .append("\nBatteryReattributeEntity{")
+                .append("\n\t" + utcToLocalTimeForLogging(timestampStart))
+                .append("\n\t" + utcToLocalTimeForLogging(timestampEnd))
+                .append("\n}");
+        return builder.toString();
+    }
+}
diff --git a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java
index 7504775..9a4f164 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/db/BatteryStateDatabase.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
 import androidx.room.Database;
 import androidx.room.Room;
 import androidx.room.RoomDatabase;
@@ -29,11 +30,13 @@
             AppUsageEventEntity.class,
             BatteryEventEntity.class,
             BatteryState.class,
-            BatteryUsageSlotEntity.class
+            BatteryUsageSlotEntity.class,
+            BatteryReattributeEntity.class
         },
         version = 1)
 public abstract class BatteryStateDatabase extends RoomDatabase {
     private static final String TAG = "BatteryStateDatabase";
+    private static final String DB_FILE_NAME = "battery-usage-db-v10";
 
     private static BatteryStateDatabase sBatteryStateDatabase;
 
@@ -49,11 +52,15 @@
     /** Provides DAO for battery usage slot table. */
     public abstract BatteryUsageSlotDao batteryUsageSlotDao();
 
+    /** Provides DAO for battery reattribution table. */
+    @NonNull
+    public abstract BatteryReattributeDao batteryReattributeDao();
+
     /** 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-v9")
+                    Room.databaseBuilder(context, BatteryStateDatabase.class, DB_FILE_NAME)
                             // Allows accessing data in the main thread for dumping bugreport.
                             .allowMainThreadQueries()
                             .fallbackToDestructiveMigration()
diff --git a/src/com/android/settings/fuelgauge/protos/battery_reattribute.proto b/src/com/android/settings/fuelgauge/protos/battery_reattribute.proto
new file mode 100644
index 0000000..8185a22
--- /dev/null
+++ b/src/com/android/settings/fuelgauge/protos/battery_reattribute.proto
@@ -0,0 +1,13 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_package = "com.android.settings.fuelgauge.batteryusage";
+option java_outer_classname = "BatteryReaatributeProto";
+
+// Battery usage reattribute data for a specific timestamp slot.
+message BatteryReattribute {
+    optional int64 timestamp_start = 1;
+    optional int64 timestamp_end = 2;
+    // Battery reattribute data for uid and its corresponding ratio.
+    map<int32, float> reattribute_data = 3;
+}
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
index b5cb446..a3b35be 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/ConvertUtilsTest.java
@@ -703,4 +703,26 @@
                                 /* taskRootPackageName= */ ""))
                 .isEqualTo(packageName);
     }
+
+    @Test
+    public void decodeBatteryReattribute_returnExpectedResult() {
+        final BatteryReattribute batteryReattribute =
+                BatteryReattribute.newBuilder()
+                        .setTimestampStart(100L)
+                        .setTimestampEnd(200L)
+                        .putReattributeData(1001, 0.2f)
+                        .putReattributeData(2001, 0.8f)
+                        .build();
+
+        final BatteryReattribute decodeResult = ConvertUtils.decodeBatteryReattribute(
+                ConvertUtils.encodeBatteryReattribute(batteryReattribute));
+
+        assertThat(decodeResult.getTimestampStart()).isEqualTo(100L);
+        assertThat(decodeResult.getTimestampEnd()).isEqualTo(200L);
+        final Map<Integer, Float> reattributeDataMap = decodeResult.getReattributeDataMap();
+        // Verify the reattribute data in the map.
+        assertThat(reattributeDataMap).hasSize(2);
+        assertThat(reattributeDataMap.get(1001)).isEqualTo(0.2f);
+        assertThat(reattributeDataMap.get(2001)).isEqualTo(0.8f);
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDaoTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDaoTest.java
new file mode 100644
index 0000000..8cb0e12
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeDaoTest.java
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2024 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 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 BatteryReattributeDao}. */
+@RunWith(RobolectricTestRunner.class)
+public final class BatteryReattributeDaoTest {
+
+    private Context mContext;
+    private BatteryStateDatabase mDatabase;
+    private BatteryReattributeDao mBatteryReattributeDao;
+
+    @Before
+    public void setUp() {
+        mContext = ApplicationProvider.getApplicationContext();
+        mDatabase = BatteryTestUtils.setUpBatteryStateDatabase(mContext);
+        mBatteryReattributeDao = mDatabase.batteryReattributeDao();
+        insert(100L, 200L, "reattributeData1");
+        insert(300L, 400L, "reattributeData3");
+        insert(200L, 300L, "reattributeData2");
+        insert(400L, 500L, "reattributeData4");
+        // Ensure there was data inserted into the database.
+        assertThat(getAllEntityData()).isNotEmpty();
+    }
+
+    @Test
+    public void getAllAfter_returnExpectedEntityData() {
+        final List<BatteryReattributeEntity> entityDataList =
+            mBatteryReattributeDao.getAllAfter(/* timestampStart= */ 300L);
+
+        assertThat(entityDataList).hasSize(2);
+        assertEntity(entityDataList.get(0), 400L, 500L, "reattributeData4");
+        assertEntity(entityDataList.get(1), 300L, 400L, "reattributeData3");
+    }
+
+    @Test
+    public void clearAll_clearAllData() {
+        mBatteryReattributeDao.clearAll();
+
+        assertThat(getAllEntityData()).isEmpty();
+    }
+
+    @Test
+    public void clearAllBefore_clearAllExpectedData() {
+        mBatteryReattributeDao.clearAllBefore(/* timestampStart= */ 300L);
+
+        final List<BatteryReattributeEntity> entityDataList = getAllEntityData();
+        assertThat(entityDataList).hasSize(1);
+        assertEntity(entityDataList.get(0), 400L, 500L, "reattributeData4");
+    }
+
+    @Test
+    public void clearAllAfter_clearAllExpectedData() {
+        mBatteryReattributeDao.clearAllAfter(/* timestampStart= */ 300L);
+
+        final List<BatteryReattributeEntity> entityDataList = getAllEntityData();
+        assertThat(entityDataList).hasSize(2);
+        assertEntity(entityDataList.get(0), 200L, 300L, "reattributeData2");
+        assertEntity(entityDataList.get(1), 100L, 200L, "reattributeData1");
+    }
+
+    @Test
+    public void insert_samePrimaryKeyEntityData_replaceIntoNewEntityData() {
+        // Verify the original data before update.
+        assertEntity(getAllEntityData().get(0), 400L, 500L, "reattributeData4");
+
+        insert(400L, 600L, "reattribute4Update");
+
+        // Verify the new update entity data.
+        assertEntity(getAllEntityData().get(0), 400L, 600L, "reattribute4Update");
+    }
+
+    private void insert(long timestampStart, long timestampEnd, String reattributeData) {
+        mBatteryReattributeDao.insert(
+                new BatteryReattributeEntity(
+                        timestampStart, timestampEnd, reattributeData));
+    }
+
+    private List<BatteryReattributeEntity> getAllEntityData() {
+        return mBatteryReattributeDao.getAllAfter(/* timestampStart= */ 0L);
+    }
+
+    private static void assertEntity(BatteryReattributeEntity entity, long timestampStart,
+            long timestampEnd, String reattributeData) {
+        assertThat(entity.timestampStart).isEqualTo(timestampStart);
+        assertThat(entity.timestampEnd).isEqualTo(timestampEnd);
+        assertThat(entity.reattributeData).isEqualTo(reattributeData);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntityTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntityTest.java
new file mode 100644
index 0000000..04912aa
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/db/BatteryReattributeEntityTest.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2024 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 com.android.settings.fuelgauge.batteryusage.BatteryReattribute;
+import com.android.settings.fuelgauge.batteryusage.ConvertUtils;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+
+/** Tests for {@link BatteryReattributeEntity}. */
+@RunWith(RobolectricTestRunner.class)
+public final class BatteryReattributeEntityTest {
+
+    @Test
+    public void constructor_createExpectedData() {
+        final BatteryReattribute batteryReattribute =
+                BatteryReattribute.newBuilder()
+                        .setTimestampStart(100L)
+                        .setTimestampEnd(200L)
+                        .putReattributeData(1001, 0.2f)
+                        .putReattributeData(2001, 0.8f)
+                        .build();
+
+        final BatteryReattributeEntity batteryReattributeEntity =
+            new BatteryReattributeEntity(batteryReattribute);
+
+        assertThat(batteryReattributeEntity.timestampStart)
+            .isEqualTo(batteryReattribute.getTimestampStart());
+        assertThat(batteryReattributeEntity.timestampEnd)
+            .isEqualTo(batteryReattribute.getTimestampEnd());
+        // Verify the BatteryReattribute data.
+        final BatteryReattribute decodeResult =
+            ConvertUtils.decodeBatteryReattribute(batteryReattributeEntity.reattributeData);
+        assertThat(decodeResult).isEqualTo(batteryReattribute);
+    }
+}