Use different context to unregisterComponentCallbacks in ViewCapture

This is required to accomadate the Launcher test suite's LeakDetector
which runs in a test rule that finishes before the ViewCaptureRule does.
This means that any context's which are stored inside ViewCapture's
normal operation will trigger a leak detection, even if there isn't one
in reality. The ViewCapture Test Rule accomadates this by deleting mRoot
before the test is finished. We can piggyback off that essentially, and
use mRoot's context to unregister the callback. This is guaranteed to
work under normal usage, and also guaranteed to not through a memory
leak detection false positive, since it will be nullified in the tests
ahead of leak detection.

Bug: 314132499
Test: Verified that post-submit didn't through a leak detection error
message.

Change-Id: I1df0dfda76f253e79b64559ce7b3111af1d37f67

2 files changed, 8 insertions, 8 deletions

diff --git a/viewcapturelib/src/com/android/app/viewcapture/NoOpViewCapture.kt b/viewcapturelib/src/com/android/app/viewcapture/NoOpViewCapture.kt
index 2b86f35..795212d 100644
--- a/viewcapturelib/src/com/android/app/viewcapture/NoOpViewCapture.kt
+++ b/viewcapturelib/src/com/android/app/viewcapture/NoOpViewCapture.kt
@@ -11,11 +11,11 @@ import android.view.Window
  */
 class NoOpViewCapture: ViewCapture(0, 0,
         createAndStartNewLooperExecutor("NoOpViewCapture", HandlerThread.MIN_PRIORITY)) {
-    override fun startCapture(view: View?, name: String?): SafeCloseable {
+    override fun startCapture(view: View, name: String): SafeCloseable {
         return SafeCloseable { }
     }
 
-    override fun startCapture(window: Window?): SafeCloseable {
+    override fun startCapture(window: Window): SafeCloseable {
         return SafeCloseable { }
     }
 }
\ No newline at end of file
diff --git a/viewcapturelib/src/com/android/app/viewcapture/ViewCapture.java b/viewcapturelib/src/com/android/app/viewcapture/ViewCapture.java
index 400bc59..6e0b047 100644
--- a/viewcapturelib/src/com/android/app/viewcapture/ViewCapture.java
+++ b/viewcapturelib/src/com/android/app/viewcapture/ViewCapture.java
@@ -128,7 +128,7 @@ public abstract class ViewCapture {
      * Attaches the ViewCapture to the provided window and returns a handle to detach the listener
      */
     @NonNull
-    public SafeCloseable startCapture(Window window) {
+    public SafeCloseable startCapture(@NonNull Window window) {
         String title = window.getAttributes().getTitle().toString();
         String name = TextUtils.isEmpty(title) ? window.toString() : title;
         return startCapture(window.getDecorView(), name);
@@ -139,16 +139,16 @@ public abstract class ViewCapture {
      * Verifies that ViewCapture is enabled before actually attaching an onDrawListener.
      */
     @NonNull
-    public SafeCloseable startCapture(View view, String name) {
+    public SafeCloseable startCapture(@NonNull View view, @NonNull String name) {
         WindowListener listener = new WindowListener(view, name);
         if (mIsEnabled) MAIN_EXECUTOR.execute(listener::attachToRoot);
         mListeners.add(listener);
-
-        Context context = view.getContext();
-        context.registerComponentCallbacks(listener);
+        view.getContext().registerComponentCallbacks(listener);
 
         return () -> {
-            context.unregisterComponentCallbacks(listener);
+            if (listener.mRoot != null && listener.mRoot.getContext() != null) {
+                listener.mRoot.getContext().unregisterComponentCallbacks(listener);
+            }
             mListeners.remove(listener);
             listener.detachFromRoot();
         };