Merge "Prevent leaked receivers from MediaOutputDialogTest" into udc-dev
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputDialogTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputDialogTest.java
index 21a7a34..425d0bc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputDialogTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputDialogTest.java
@@ -39,6 +39,7 @@
 import android.util.FeatureFlagUtils;
 import android.view.View;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.MediumTest;
 
 import com.android.internal.logging.UiEventLogger;
@@ -64,6 +65,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Consumer;
 
 @MediumTest
 @RunWith(AndroidTestingRunner.class)
@@ -89,7 +91,7 @@
     private final UiEventLogger mUiEventLogger = mock(UiEventLogger.class);
     private final DialogLaunchAnimator mDialogLaunchAnimator = mock(DialogLaunchAnimator.class);
     private final MediaMetadata mMediaMetadata = mock(MediaMetadata.class);
-    private final MediaDescription  mMediaDescription = mock(MediaDescription.class);
+    private final MediaDescription mMediaDescription = mock(MediaDescription.class);
     private final NearbyMediaDevicesManager mNearbyMediaDevicesManager = mock(
             NearbyMediaDevicesManager.class);
     private final AudioManager mAudioManager = mock(AudioManager.class);
@@ -102,6 +104,11 @@
     private MediaOutputController mMediaOutputController;
     private final List<String> mFeatures = new ArrayList<>();
 
+    @Override
+    protected boolean shouldFailOnLeakedReceiver() {
+        return true;
+    }
+
     @Before
     public void setUp() {
         when(mLocalBluetoothManager.getProfileManager()).thenReturn(mLocalBluetoothProfileManager);
@@ -120,8 +127,7 @@
                 Optional.of(mNearbyMediaDevicesManager), mAudioManager, mPowerExemptionManager,
                 mKeyguardManager, mFlags);
         mMediaOutputController.mLocalMediaManager = mLocalMediaManager;
-        mMediaOutputDialog = new MediaOutputDialog(mContext, false, mBroadcastSender,
-                mMediaOutputController, mUiEventLogger);
+        mMediaOutputDialog = makeTestDialog(mMediaOutputController);
         mMediaOutputDialog.show();
 
         when(mLocalMediaManager.getCurrentConnectedDevice()).thenReturn(mMediaDevice);
@@ -130,7 +136,7 @@
 
     @After
     public void tearDown() {
-        mMediaOutputDialog.dismissDialog();
+        mMediaOutputDialog.dismiss();
     }
 
     @Test
@@ -311,11 +317,9 @@
         MediaOutputController mockMediaOutputController = mock(MediaOutputController.class);
         when(mockMediaOutputController.isBroadcastSupported()).thenReturn(false);
 
-        MediaOutputDialog testDialog = new MediaOutputDialog(mContext, false, mBroadcastSender,
-                mockMediaOutputController, mUiEventLogger);
-        testDialog.show();
-
-        assertThat(testDialog.getStopButtonText().toString()).isEqualTo(stopText);
+        withTestDialog(mockMediaOutputController, testDialog -> {
+            assertThat(testDialog.getStopButtonText().toString()).isEqualTo(stopText);
+        });
     }
 
     @Test
@@ -328,11 +332,9 @@
         when(mockMediaOutputController.isBluetoothLeDevice(any())).thenReturn(true);
         when(mockMediaOutputController.isPlaying()).thenReturn(true);
         when(mockMediaOutputController.isBluetoothLeBroadcastEnabled()).thenReturn(false);
-        MediaOutputDialog testDialog = new MediaOutputDialog(mContext, false, mBroadcastSender,
-                mockMediaOutputController, mUiEventLogger);
-        testDialog.show();
-
-        assertThat(testDialog.getStopButtonText().toString()).isEqualTo(stopText);
+        withTestDialog(mockMediaOutputController, testDialog -> {
+            assertThat(testDialog.getStopButtonText().toString()).isEqualTo(stopText);
+        });
     }
 
     @Test
@@ -341,11 +343,9 @@
         when(mockMediaOutputController.isBroadcastSupported()).thenReturn(false);
         when(mockMediaOutputController.getCurrentConnectedMediaDevice()).thenReturn(null);
         when(mockMediaOutputController.isPlaying()).thenReturn(false);
-        MediaOutputDialog testDialog = new MediaOutputDialog(mContext, false, mBroadcastSender,
-                mockMediaOutputController, mUiEventLogger);
-        testDialog.show();
-
-        testDialog.onStopButtonClick();
+        withTestDialog(mockMediaOutputController, testDialog -> {
+            testDialog.onStopButtonClick();
+        });
 
         verify(mockMediaOutputController).releaseSession();
     }
@@ -354,13 +354,22 @@
     // Check the visibility metric logging by creating a new MediaOutput dialog,
     // and verify if the calling times increases.
     public void onCreate_ShouldLogVisibility() {
-        MediaOutputDialog testDialog = new MediaOutputDialog(mContext, false, mBroadcastSender,
-                mMediaOutputController, mUiEventLogger);
-        testDialog.show();
-
-        testDialog.dismissDialog();
+        withTestDialog(mMediaOutputController, testDialog -> {});
 
         verify(mUiEventLogger, times(2))
                 .log(MediaOutputDialog.MediaOutputEvent.MEDIA_OUTPUT_DIALOG_SHOW);
     }
+
+    @NonNull
+    private MediaOutputDialog makeTestDialog(MediaOutputController controller) {
+        return new MediaOutputDialog(mContext, false, mBroadcastSender,
+                controller, mUiEventLogger);
+    }
+
+    private void withTestDialog(MediaOutputController controller, Consumer<MediaOutputDialog> c) {
+        MediaOutputDialog testDialog = makeTestDialog(controller);
+        testDialog.show();
+        c.accept(testDialog);
+        testDialog.dismiss();
+    }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java b/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
index 1ec4e8c..8bbd58d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
@@ -102,7 +102,8 @@
                 mock(Executor.class),
                 mock(DumpManager.class),
                 mock(BroadcastDispatcherLogger.class),
-                mock(UserTracker.class));
+                mock(UserTracker.class),
+                shouldFailOnLeakedReceiver());
 
         mRealInstrumentation = InstrumentationRegistry.getInstrumentation();
         Instrumentation inst = spy(mRealInstrumentation);
@@ -141,6 +142,10 @@
         mDependency.injectTestDependency(DialogLaunchAnimator.class, fakeDialogLaunchAnimator());
     }
 
+    protected boolean shouldFailOnLeakedReceiver() {
+        return false;
+    }
+
     @After
     public void SysuiTeardown() {
         if (mRealInstrumentation != null) {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
index ad086ff..af940e4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.broadcast.logging.BroadcastDispatcherLogger
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.settings.UserTracker
+import java.lang.IllegalStateException
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.Executor
 
@@ -37,7 +38,8 @@
     broadcastRunningExecutor: Executor,
     dumpManager: DumpManager,
     logger: BroadcastDispatcherLogger,
-    userTracker: UserTracker
+    userTracker: UserTracker,
+    private val shouldFailOnLeakedReceiver: Boolean
 ) :
     BroadcastDispatcher(
         context,
@@ -85,6 +87,9 @@
     fun cleanUpReceivers(testName: String) {
         registeredReceivers.forEach {
             Log.i(testName, "Receiver not unregistered from dispatcher: $it")
+            if (shouldFailOnLeakedReceiver) {
+                throw IllegalStateException("Receiver not unregistered from dispatcher: $it")
+            }
         }
         registeredReceivers.clear()
     }