diff --git a/core/java/android/os/TEST_MAPPING b/core/java/android/os/TEST_MAPPING
index 97e03e9d0d94b..55b1f940c7644 100644
--- a/core/java/android/os/TEST_MAPPING
+++ b/core/java/android/os/TEST_MAPPING
@@ -40,7 +40,12 @@
       ]
     },
     {
-      "file_patterns": ["BatteryStats.java"],
+      "file_patterns": [
+        "BatteryStats[^/]*\\.java",
+        "BatteryUsageStats[^/]*\\.java",
+        "PowerComponents\\.java",
+        "[^/]*BatteryConsumer[^/]*\\.java"
+      ],
       "name": "FrameworksCoreTests",
       "options": [
         { "include-filter": "com.android.internal.os.BatteryStatsTests" },
@@ -48,13 +53,26 @@
       ]
     },
     {
-      "file_patterns": ["BatteryStats.java"],
+      "file_patterns": [
+        "BatteryStats[^/]*\\.java",
+        "BatteryUsageStats[^/]*\\.java",
+        "PowerComponents\\.java",
+        "[^/]*BatteryConsumer[^/]*\\.java"
+      ],
       "name": "FrameworksServicesTests",
       "options": [
         { "include-filter": "com.android.server.am.BatteryStatsServiceTest" },
         { "include-filter": "com.android.server.am.MeasuredEnergySnapshotTest" },
         { "include-filter": "com.android.server.am.BatteryExternalStatsWorkerTest" }
       ]
+    },
+    {
+      "file_patterns": [
+        "BatteryUsageStats[^/]*\\.java",
+        "PowerComponents\\.java",
+        "[^/]*BatteryConsumer[^/]*\\.java"
+      ],
+      "name": "BatteryUsageStatsProtoTests"
     }
   ],
   "postsubmit": [
diff --git a/core/java/com/android/internal/os/TEST_MAPPING b/core/java/com/android/internal/os/TEST_MAPPING
index 2b22f08652e73..5a5165ddd7e2e 100644
--- a/core/java/com/android/internal/os/TEST_MAPPING
+++ b/core/java/com/android/internal/os/TEST_MAPPING
@@ -12,6 +12,14 @@
         { "exclude-annotation": "com.android.internal.os.SkipPresubmit" }
       ]
     },
+    {
+      "file_patterns": [
+        "Battery[^/]*\\.java",
+        "Kernel[^/]*\\.java",
+        "[^/]*Power[^/]*\\.java"
+      ],
+      "name": "BatteryUsageStatsProtoTests"
+    },
     {
       "file_patterns": [
         "BinderDeathDispatcher\\.java"
@@ -23,7 +31,11 @@
       ]
     },
     {
-      "file_patterns": ["Battery[^/]*\\.java"],
+      "file_patterns": [
+        "Battery[^/]*\\.java",
+        "Kernel[^/]*\\.java",
+        "[^/]*Power[^/]*\\.java"
+      ],
       "name": "FrameworksServicesTests",
       "options": [
         { "include-filter": "com.android.server.am.BatteryStatsServiceTest" },