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

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);
         }

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);

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);
 

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));

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) {