Add warnings for leaked receivers in tests

This CL adds a way to track BroadcastReceivers that are leaking in test
cases that use SysuiTestCase (all of SystemUITests). Every receiver that
is registered with Context is recorded and those that are not
unregistered are logged as a Warning. After each test, receivers are
unregistered so they don't polute following tests.

Additionally, if a BroadcastDispatcher is instantiated, a Fake is
provided that will also track leaked receivers and log them as Info (as
BroadcastDispatcher will only keep WeakReferences).

This is a separate tracker done directly in SysuiTestableContext instead
of using LeakCheckedTest or a LeakChecker as that is usually pretty
stringent (it will fail the test).

Test: SystemUITests
Bug: 151614195
Change-Id: I11afb49ce3bbc51d98fa069ba4074c852d07bcce
diff --git a/packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt b/packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt
index bd803fa..e89c66e 100644
--- a/packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt
+++ b/packages/SystemUI/src/com/android/systemui/broadcast/BroadcastDispatcher.kt
@@ -94,7 +94,7 @@
     @Deprecated(message = "Replacing Handler for Executor in SystemUI",
             replaceWith = ReplaceWith("registerReceiver(receiver, filter, executor, user)"))
     @JvmOverloads
-    fun registerReceiverWithHandler(
+    open fun registerReceiverWithHandler(
         receiver: BroadcastReceiver,
         filter: IntentFilter,
         handler: Handler,
@@ -118,7 +118,7 @@
      *                                  categories or the filter has no actions.
      */
     @JvmOverloads
-    fun registerReceiver(
+    open fun registerReceiver(
         receiver: BroadcastReceiver,
         filter: IntentFilter,
         executor: Executor? = context.mainExecutor,
@@ -149,7 +149,7 @@
      *
      * @param receiver The receiver to unregister. It will be unregistered for all users.
      */
-    fun unregisterReceiver(receiver: BroadcastReceiver) {
+    open fun unregisterReceiver(receiver: BroadcastReceiver) {
         handler.obtainMessage(MSG_REMOVE_RECEIVER, receiver).sendToTarget()
     }
 
@@ -159,7 +159,7 @@
      * @param receiver The receiver to unregister. It will be unregistered for all users.
      * @param user The user associated to the registered [receiver]. It can be [UserHandle.ALL].
      */
-    fun unregisterReceiverForUser(receiver: BroadcastReceiver, user: UserHandle) {
+    open fun unregisterReceiverForUser(receiver: BroadcastReceiver, user: UserHandle) {
         handler.obtainMessage(MSG_REMOVE_RECEIVER_FOR_USER, user.identifier, 0, receiver)
                 .sendToTarget()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
index bb2eea9..32604f8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestCase.java
@@ -15,6 +15,7 @@
  */
 package com.android.systemui;
 
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
@@ -32,7 +33,10 @@
 
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.settingslib.bluetooth.LocalBluetoothManager;
+import com.android.systemui.broadcast.BroadcastDispatcher;
+import com.android.systemui.broadcast.FakeBroadcastDispatcher;
 import com.android.systemui.classifier.FalsingManagerFake;
+import com.android.systemui.dump.DumpManager;
 import com.android.systemui.plugins.FalsingManager;
 
 import org.junit.After;
@@ -60,11 +64,14 @@
             new DexmakerShareClassLoaderRule();
     public TestableDependency mDependency;
     private Instrumentation mRealInstrumentation;
+    private FakeBroadcastDispatcher mFakeBroadcastDispatcher;
 
     @Before
     public void SysuiSetup() throws Exception {
         SystemUIFactory.createFromConfig(mContext);
         mDependency = new TestableDependency(mContext);
+        mFakeBroadcastDispatcher = new FakeBroadcastDispatcher(mContext, mock(Handler.class),
+                mock(Looper.class), mock(DumpManager.class));
 
         mRealInstrumentation = InstrumentationRegistry.getInstrumentation();
         Instrumentation inst = spy(mRealInstrumentation);
@@ -77,12 +84,18 @@
                     "SysUI Tests should use SysuiTestCase#getContext or SysuiTestCase#mContext");
         });
         InstrumentationRegistry.registerInstance(inst, InstrumentationRegistry.getArguments());
+        // Many tests end up creating a BroadcastDispatcher. Instead, give them a fake that will
+        // record receivers registered. They are not actually leaked as they are kept just as a weak
+        // reference and are never sent to the Context. This will also prevent a real
+        // BroadcastDispatcher from actually registering receivers.
+        mDependency.injectTestDependency(BroadcastDispatcher.class, mFakeBroadcastDispatcher);
         // A lot of tests get the FalsingManager, often via several layers of indirection.
         // None of them actually need it.
         mDependency.injectTestDependency(FalsingManager.class, new FalsingManagerFake());
         mDependency.injectMockDependency(KeyguardUpdateMonitor.class);
 
-        // TODO: b/151614195 investigate root cause of needing this mock dependency
+        // A lot of tests get the LocalBluetoothManager, often via several layers of indirection.
+        // None of them actually need it.
         mDependency.injectMockDependency(LocalBluetoothManager.class);
     }
 
@@ -95,6 +108,8 @@
         }
         disallowTestableLooperAsMainThread();
         SystemUIFactory.cleanup();
+        mContext.cleanUpReceivers(this.getClass().getSimpleName());
+        mFakeBroadcastDispatcher.cleanUpReceivers(this.getClass().getSimpleName());
     }
 
     /**
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestableContext.java b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestableContext.java
index d847208..95ff98a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/SysuiTestableContext.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/SysuiTestableContext.java
@@ -14,12 +14,24 @@
 
 package com.android.systemui;
 
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.os.Handler;
+import android.os.UserHandle;
 import android.testing.LeakCheck;
 import android.testing.TestableContext;
+import android.util.ArraySet;
+import android.util.Log;
 import android.view.Display;
 
+import java.util.Set;
+
 public class SysuiTestableContext extends TestableContext {
+
+    private Set<BroadcastReceiver> mRegisteredReceivers = new ArraySet<>();
+
     public SysuiTestableContext(Context base) {
         super(base);
         setTheme(R.style.Theme_SystemUI);
@@ -40,4 +52,42 @@
                 new SysuiTestableContext(getBaseContext().createDisplayContext(display));
         return context;
     }
+
+    public void cleanUpReceivers(String testName) {
+        Set<BroadcastReceiver> copy = new ArraySet<>(mRegisteredReceivers);
+        for (BroadcastReceiver r : copy) {
+            try {
+                unregisterReceiver(r);
+                Log.w(testName, "Receiver not unregistered from Context: " + r);
+            } catch (IllegalArgumentException e) {
+                // Nothing to do here. Somehow it got unregistered.
+            }
+        }
+    }
+
+    @Override
+    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter) {
+        mRegisteredReceivers.add(receiver);
+        return super.registerReceiver(receiver, filter);
+    }
+
+    @Override
+    public Intent registerReceiver(BroadcastReceiver receiver, IntentFilter filter,
+            String broadcastPermission, Handler scheduler) {
+        mRegisteredReceivers.add(receiver);
+        return super.registerReceiver(receiver, filter, broadcastPermission, scheduler);
+    }
+
+    @Override
+    public Intent registerReceiverAsUser(BroadcastReceiver receiver, UserHandle user,
+            IntentFilter filter, String broadcastPermission, Handler scheduler) {
+        mRegisteredReceivers.add(receiver);
+        return super.registerReceiverAsUser(receiver, user, filter, broadcastPermission, scheduler);
+    }
+
+    @Override
+    public void unregisterReceiver(BroadcastReceiver receiver) {
+        mRegisteredReceivers.remove(receiver);
+        super.unregisterReceiver(receiver);
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt b/packages/SystemUI/tests/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
new file mode 100644
index 0000000..9a5773a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/broadcast/FakeBroadcastDispatcher.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.broadcast
+
+import android.content.BroadcastReceiver
+import android.content.IntentFilter
+import android.os.Handler
+import android.os.Looper
+import android.os.UserHandle
+import android.util.ArraySet
+import android.util.Log
+import com.android.systemui.SysuiTestableContext
+import com.android.systemui.dump.DumpManager
+import java.util.concurrent.Executor
+
+class FakeBroadcastDispatcher(
+    context: SysuiTestableContext,
+    handler: Handler,
+    looper: Looper,
+    dumpManager: DumpManager
+) : BroadcastDispatcher(context, handler, looper, dumpManager) {
+
+    private val registeredReceivers = ArraySet<BroadcastReceiver>()
+
+    override fun registerReceiverWithHandler(
+        receiver: BroadcastReceiver,
+        filter: IntentFilter,
+        handler: Handler,
+        user: UserHandle
+    ) {
+        registeredReceivers.add(receiver)
+    }
+
+    override fun registerReceiver(
+        receiver: BroadcastReceiver,
+        filter: IntentFilter,
+        executor: Executor?,
+        user: UserHandle
+    ) {
+        registeredReceivers.add(receiver)
+    }
+
+    override fun unregisterReceiver(receiver: BroadcastReceiver) {
+        registeredReceivers.remove(receiver)
+    }
+
+    override fun unregisterReceiverForUser(receiver: BroadcastReceiver, user: UserHandle) {
+        registeredReceivers.remove(receiver)
+    }
+
+    fun cleanUpReceivers(testName: String) {
+        registeredReceivers.forEach {
+            Log.i(testName, "Receiver not unregistered from dispatcher: $it")
+        }
+        registeredReceivers.clear()
+    }
+}
\ No newline at end of file