Fix several flaky DeviceStateManagerServiceTests

Fixes: 297949293
Bug: 223153452

Test: run each test with 100 iterations
Change-Id: I333dcc2c10346ce585d1c5e7f16cad113b2b5a95
Merged-In: I333dcc2c10346ce585d1c5e7f16cad113b2b5a95

1 files changed, 45 insertions, 54 deletions

diff --git a/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
index 94f88abf7301..567410145831 100644
--- a/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
@@ -18,6 +18,8 @@ package com.android.server.devicestate;
 
 import static android.hardware.devicestate.DeviceStateManager.INVALID_DEVICE_STATE;
 
+import static com.android.compatibility.common.util.PollingCheck.waitFor;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.mock;
@@ -40,6 +42,7 @@ import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.FlakyTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.compatibility.common.util.PollingCheck;
 import com.android.server.wm.ActivityTaskManagerInternal;
 import com.android.server.wm.WindowProcessController;
 
@@ -73,6 +76,8 @@ public final class DeviceStateManagerServiceTest {
 
     private static final int FAKE_PROCESS_ID = 100;
 
+    private static final int TIMEOUT = 2000;
+
     private TestDeviceStatePolicy mPolicy;
     private TestDeviceStateProvider mProvider;
     private DeviceStateManagerService mService;
@@ -106,6 +111,10 @@ public final class DeviceStateManagerServiceTest {
         }
     }
 
+    private void waitAndAssert(PollingCheck.PollingCheckCondition condition) {
+        waitFor(TIMEOUT, condition);
+    }
+
     @Test
     public void baseStateChanged() {
         assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
@@ -272,37 +281,34 @@ public final class DeviceStateManagerServiceTest {
         mService.getBinderService().registerCallback(callback);
 
         mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
-        flushHandler();
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                OTHER_DEVICE_STATE.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState,
-                OTHER_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == OTHER_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().currentState
+                == OTHER_DEVICE_STATE.getIdentifier());
 
         mProvider.setState(DEFAULT_DEVICE_STATE.getIdentifier());
-        flushHandler();
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                DEFAULT_DEVICE_STATE.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState,
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == DEFAULT_DEVICE_STATE.getIdentifier());
+
+        waitAndAssert(() -> callback.getLastNotifiedInfo().currentState
+                == DEFAULT_DEVICE_STATE.getIdentifier());
 
         mPolicy.blockConfigure();
         mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
-        flushHandler();
         // The callback should not have been notified of the state change as the policy is still
         // pending callback.
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                DEFAULT_DEVICE_STATE.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState,
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == DEFAULT_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().currentState
+                == DEFAULT_DEVICE_STATE.getIdentifier());
 
         mPolicy.resumeConfigure();
-        flushHandler();
         // Now that the policy is finished processing the callback should be notified of the state
         // change.
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                OTHER_DEVICE_STATE.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState,
-                OTHER_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == OTHER_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().currentState
+                == OTHER_DEVICE_STATE.getIdentifier());
     }
 
     @Test
@@ -329,13 +335,9 @@ public final class DeviceStateManagerServiceTest {
 
         mService.getBinderService().requestState(token, OTHER_DEVICE_STATE.getIdentifier(),
                 0 /* flags */);
-        // Flush the handler twice. The first flush ensures the request is added and the policy is
-        // notified, while the second flush ensures the callback is notified once the change is
-        // committed.
-        flushHandler(2 /* count */);
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
         assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
         assertEquals(mSysPropSetter.getValue(),
@@ -352,12 +354,11 @@ public final class DeviceStateManagerServiceTest {
                 OTHER_DEVICE_STATE.getIdentifier());
 
         mService.getBinderService().cancelStateRequest();
-        flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state once the override is cleared.
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
+        waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(DEFAULT_DEVICE_STATE)));
         assertEquals(mSysPropSetter.getValue(),
                 DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
         assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
@@ -600,13 +601,9 @@ public final class DeviceStateManagerServiceTest {
         mService.getBinderService().requestBaseStateOverride(token,
                 OTHER_DEVICE_STATE.getIdentifier(),
                 0 /* flags */);
-        // Flush the handler twice. The first flush ensures the request is added and the policy is
-        // notified, while the second flush ensures the callback is notified once the change is
-        // committed.
-        flushHandler(2 /* count */);
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
         assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
         assertEquals(mSysPropSetter.getValue(),
@@ -624,12 +621,11 @@ public final class DeviceStateManagerServiceTest {
                 OTHER_DEVICE_STATE.getIdentifier());
 
         mService.getBinderService().cancelBaseStateOverride();
-        flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state once the override is cleared.
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
+        waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(DEFAULT_DEVICE_STATE)));
         assertEquals(mSysPropSetter.getValue(),
                 DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
         assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
@@ -638,8 +634,8 @@ public final class DeviceStateManagerServiceTest {
         assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
                 DEFAULT_DEVICE_STATE.getIdentifier());
 
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == DEFAULT_DEVICE_STATE.getIdentifier());
         assertEquals(callback.getLastNotifiedInfo().currentState,
                 DEFAULT_DEVICE_STATE.getIdentifier());
     }
@@ -660,13 +656,9 @@ public final class DeviceStateManagerServiceTest {
         mService.getBinderService().requestBaseStateOverride(token,
                 OTHER_DEVICE_STATE.getIdentifier(),
                 0 /* flags */);
-        // Flush the handler twice. The first flush ensures the request is added and the policy is
-        // notified, while the second flush ensures the callback is notified once the change is
-        // committed.
-        flushHandler(2 /* count */);
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
         assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
         assertEquals(mSysPropSetter.getValue(),
@@ -684,12 +676,11 @@ public final class DeviceStateManagerServiceTest {
                 OTHER_DEVICE_STATE.getIdentifier());
 
         mProvider.setState(testDeviceState.getIdentifier());
-        flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        waitAndAssert(() -> callback.getLastNotifiedStatus(token)
+                == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set to the new base state once the override is cleared.
-        assertEquals(mService.getCommittedState(), Optional.of(testDeviceState));
+        waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(testDeviceState)));
         assertEquals(mSysPropSetter.getValue(),
                 testDeviceState.getIdentifier() + ":" + testDeviceState.getName());
         assertEquals(mService.getBaseState(), Optional.of(testDeviceState));
@@ -698,8 +689,8 @@ public final class DeviceStateManagerServiceTest {
         assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
                 testDeviceState.getIdentifier());
 
-        assertEquals(callback.getLastNotifiedInfo().baseState,
-                testDeviceState.getIdentifier());
+        waitAndAssert(() -> callback.getLastNotifiedInfo().baseState
+                == testDeviceState.getIdentifier());
         assertEquals(callback.getLastNotifiedInfo().currentState,
                 testDeviceState.getIdentifier());
     }