Reduce the chance of BatteryStats background work affecting CTS tests

Add some utility methods to help get consistent behavior from
batterystats

Fixes: 185970842
Test: atest
android.cts.statsd.validation.ValidationTests#testPartialWakelock
--rerun-until-failure

Change-Id: I773a4afb91822103a8a666b575f128d776bed5ec
diff --git a/tests/src/android/cts/statsd/atom/AtomTestCase.java b/tests/src/android/cts/statsd/atom/AtomTestCase.java
index c7e984e..5f3d902 100644
--- a/tests/src/android/cts/statsd/atom/AtomTestCase.java
+++ b/tests/src/android/cts/statsd/atom/AtomTestCase.java
@@ -993,6 +993,19 @@
         getDevice().executeShellCommand("cmd battery reset");
     }
 
+    protected void turnBatteryStatsAutoResetOn() throws Exception {
+        getDevice().executeShellCommand("dumpsys batterystats enable no-auto-reset");
+    }
+
+    protected void turnBatteryStatsAutoResetOff() throws Exception {
+        getDevice().executeShellCommand("dumpsys batterystats enable no-auto-reset");
+    }
+
+    protected void flushBatteryStatsHandlers() throws Exception {
+        // Dumping batterystats will flush everything in the batterystats handler threads.
+        getDevice().executeShellCommand(DUMP_BATTERYSTATS_CMD);
+    }
+
     protected void rebootDevice() throws Exception {
         getDevice().rebootUntilOnline();
     }
diff --git a/tests/src/android/cts/statsd/validation/ValidationTests.java b/tests/src/android/cts/statsd/validation/ValidationTests.java
index 6a3fa6a..90b8fb4 100644
--- a/tests/src/android/cts/statsd/validation/ValidationTests.java
+++ b/tests/src/android/cts/statsd/validation/ValidationTests.java
@@ -69,12 +69,14 @@
     @Override
     protected void setUp() throws Exception {
         super.setUp();
+        turnBatteryStatsAutoResetOff(); // Turn off Battery Stats auto resetting
     }
 
     @Override
     protected void tearDown() throws Exception {
         resetBatteryStatus(); // Undo any unplugDevice().
         turnScreenOn(); // Reset screen to on state
+        turnBatteryStatsAutoResetOn(); // Turn Battery Stats auto resetting back on
         super.tearDown();
     }
 
@@ -82,6 +84,7 @@
         if (!hasFeature(FEATURE_AUTOMOTIVE, false)) return;
         resetBatteryStats();
         unplugDevice();
+        flushBatteryStatsHandlers();
         // AoD needs to be turned off because the screen should go into an off state. But, if AoD is
         // on and the device doesn't support STATE_DOZE, the screen sadly goes back to STATE_ON.
         String aodState = getAodState();