Reduce some flaky cases of WmTests

1. DisplayPolicy -> ScreenshotHelper register real broadcast
   receiver because there are other overloading
   registerReceiver methods that are not mocked.
   If the number exceeds 1000, it will crash.
2. The transition logger may access mocked objects on another
   thread, which will crash the test.
3. Wait for activity operations to complete before leaving
   test that releases shell permission. Otherwise the crash
   (no INTERNAL_SYSTEM_WINDOW) on main thread will stop the test.

Bug: 254268165
Test: WmTests
Change-Id: Ie2a7ac74724ff38245641f43338010adb5ed6ea8

5 files changed, 27 insertions, 5 deletions

diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 1b27bb17f599..814bd85f553f 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1603,7 +1603,7 @@ class Transition implements BLASTSyncEngine.TransactionReadyListener {
         // Since we created root-leash but no longer reference it from core, release it now
         info.releaseAnimSurfaces();
 
-        mController.mLoggerHandler.post(mLogger::logOnSend);
+        mLogger.logOnSendAsync(mController.mLoggerHandler);
         if (mLogger.mInfo != null) {
             mController.mTransitionTracer.logSentTransition(this, mTargets);
         }
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index 79cb61be5948..881eddc03243 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -1434,7 +1434,7 @@ class TransitionController {
      * Beside `mCreateWallTimeMs`, all times are elapsed times and will all be reported relative
      * to when the transition was created.
      */
-    static class Logger {
+    static class Logger implements Runnable {
         long mCreateWallTimeMs;
         long mCreateTimeNs;
         long mRequestTimeNs;
@@ -1458,6 +1458,20 @@ class TransitionController {
             return sb.toString();
         }
 
+        void logOnSendAsync(Handler handler) {
+            handler.post(this);
+        }
+
+        @Override
+        public void run() {
+            try {
+                logOnSend();
+            } catch (Exception e) {
+                // In case TransitionRequestInfo#toString() accesses window container with race.
+                Slog.w(TAG, "Failed to log transition", e);
+            }
+        }
+
         void logOnSend() {
             ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS_MIN, "%s", buildOnSendLog());
             ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS_MIN, "    startWCT=%s", mStartWCT);
diff --git a/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java b/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
index 41bfc806f839..1a4b94b074a1 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
@@ -98,7 +98,7 @@ public class SurfaceControlViewHostTests {
         mView1 = new Button(mActivity);
         mView2 = new Button(mActivity);
 
-        mActivity.runOnUiThread(() -> {
+        mInstrumentation.runOnMainSync(() -> {
             TestWindowlessWindowManager wwm = new TestWindowlessWindowManager(
                     mActivity.getResources().getConfiguration(), sc, null);
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 7edfd9a3ecb3..e252b9e1c8be 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -208,7 +208,11 @@ public class SystemServicesTestRule implements TestRule {
         spyOn(mContext);
 
         doReturn(null).when(mContext)
-                .registerReceiver(nullable(BroadcastReceiver.class), any(IntentFilter.class));
+                .registerReceiver(nullable(BroadcastReceiver.class), any(IntentFilter.class),
+                        nullable(String.class), nullable(Handler.class));
+        doReturn(null).when(mContext)
+                .registerReceiver(nullable(BroadcastReceiver.class), any(IntentFilter.class),
+                        nullable(String.class), nullable(Handler.class), anyInt());
         doReturn(null).when(mContext)
                 .registerReceiverAsUser(any(BroadcastReceiver.class), any(UserHandle.class),
                         any(IntentFilter.class), nullable(String.class), nullable(Handler.class));
diff --git a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
index ed0c8ef489e5..7351e81387bb 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -114,7 +114,11 @@ public class TransitionTests extends WindowTestsBase {
     private BLASTSyncEngine mSyncEngine;
 
     private Transition createTestTransition(int transitType, TransitionController controller) {
-        return new Transition(transitType, 0 /* flags */, controller, controller.mSyncEngine);
+        final Transition transition = new Transition(transitType, 0 /* flags */, controller,
+                controller.mSyncEngine);
+        spyOn(transition.mLogger);
+        doNothing().when(transition.mLogger).logOnSendAsync(any());
+        return transition;
     }
 
     private Transition createTestTransition(int transitType) {