Merge "Fix several flaky DeviceStateManagerServiceTests" into main

1 files changed, 45 insertions, 55 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 668415041129..8f88ba6ff9a7 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;
@@ -41,6 +43,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;
 
@@ -81,6 +84,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;
@@ -118,6 +123,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));
@@ -291,44 +300,40 @@ public final class DeviceStateManagerServiceTest {
         assertEquals(info.currentState, INVALID_DEVICE_STATE);
     }
 
-    @FlakyTest(bugId = 223153452)
     @Test
     public void registerCallback() throws RemoteException {
         TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         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
@@ -370,13 +375,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(),
@@ -393,12 +394,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));
@@ -691,13 +691,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(),
@@ -715,12 +711,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));
@@ -729,8 +724,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());
     }
@@ -751,13 +746,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(),
@@ -775,12 +766,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));
@@ -789,8 +779,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());
     }