Merge "Fix flaky AppWindowTokenTests"
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index bcf6aba..727fb3b 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -1657,8 +1657,10 @@
             final ActivityManager.TaskSnapshot snapshot = snapshotCtrl.getSnapshot(
                     getTask().mTaskId, getTask().mUserId, false /* restoreFromDisk */,
                     false /* reducedResolution */);
-            mThumbnail = new AppWindowThumbnail(t, this, snapshot.getSnapshot(),
-                    true /* relative */);
+            if (snapshot != null) {
+                mThumbnail = new AppWindowThumbnail(t, this, snapshot.getSnapshot(),
+                        true /* relative */);
+            }
         }
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
index 069a609..e5e850a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppWindowTokenTests.java
@@ -62,7 +62,6 @@
  * Build/Install/Run:
  *  atest FrameworksServicesTests:AppWindowTokenTests
  */
-@FlakyTest(bugId = 68267650)
 @SmallTest
 @Presubmit
 public class AppWindowTokenTests extends WindowTestsBase {
@@ -79,6 +78,7 @@
         mTask = createTaskInStack(mStack, 0 /* userId */);
         mToken = WindowTestUtils.createTestAppWindowToken(mDisplayContent);
 
+        mToken.mSkipOnParentChanged = false;
         mTask.addChild(mToken, 0);
     }
 
@@ -303,6 +303,7 @@
         assertNoStartingWindow(mToken);
     }
 
+    @FlakyTest(detail = "Promote to presubmit when shown to be stable.")
     @Test
     public void testAddRemoveRace() {
         // There was once a race condition between adding and removing starting windows
@@ -387,6 +388,7 @@
         // bottom one.
         tokenTop.setVisibility(false, false);
         tokenBottom.transferStartingWindowFromHiddenAboveTokenIfNeeded();
+        waitUntilHandlersIdle();
 
         // Assert that the bottom window now has the starting window.
         assertNoStartingWindow(tokenTop);
diff --git a/services/tests/wmtests/src/com/android/server/wm/TestSystemServices.java b/services/tests/wmtests/src/com/android/server/wm/TestSystemServices.java
index 1c1fe29..e540b3a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TestSystemServices.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TestSystemServices.java
@@ -61,7 +61,7 @@
 import org.mockito.invocation.InvocationOnMock;
 import org.mockito.quality.Strictness;
 
-import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * A Test utility class to create a mock {@link WindowManagerService} instance for tests.
@@ -71,6 +71,8 @@
     private static WindowManagerService sService;
     private static TestWindowManagerPolicy sPolicy;
 
+    static AtomicBoolean sCurrentMessagesProcessed = new AtomicBoolean(false);
+
     static void setUpWindowManagerService() {
         sMockitoSession = mockitoSession()
                 .spyStatic(LockGuard.class)
@@ -195,21 +197,23 @@
     }
 
     private static void waitHandlerIdle(Handler handler) {
-        if (!handler.hasMessagesOrCallbacks()) {
-            return;
-        }
-        final CountDownLatch latch = new CountDownLatch(1);
-        // Wait for delayed messages are processed.
-        handler.getLooper().getQueue().addIdleHandler(() -> {
-            if (handler.hasMessagesOrCallbacks()) {
-                return true; // keep idle handler.
+        synchronized (sCurrentMessagesProcessed) {
+            // Add a message to the handler queue and make sure it is fully processed before we move
+            // on. This makes sure all previous messages in the handler are fully processed vs. just
+            // popping them from the message queue.
+            sCurrentMessagesProcessed.set(false);
+            handler.post(() -> {
+                synchronized (sCurrentMessagesProcessed) {
+                    sCurrentMessagesProcessed.set(true);
+                    sCurrentMessagesProcessed.notifyAll();
+                }
+            });
+            while (!sCurrentMessagesProcessed.get()) {
+                try {
+                    sCurrentMessagesProcessed.wait();
+                } catch (InterruptedException e) {
+                }
             }
-            latch.countDown();
-            return false; // remove idle handler.
-        });
-        try {
-            latch.await();
-        } catch (InterruptedException e) {
         }
     }
 }