repeatWhenAttached

This new View extension function replaces WindowAddedViewLifecycleOwner and is
intended for use with views that are not part of an activity.

It is more correct because it properly disposes itself and stops
all previously launched coroutines/jobs when the view is detached from
its view hierarchy.

Test: Extensive unit tests included. Also tested manually making sure that there are no
crashes and that jobs scheduled by a view-binder are properly cleaned up
when the view is detached and replaced by a different view when changing
device configuration using:

$ adb shell wm density 1000

Bug: 235403546
Change-Id: Ied36c9e1735333c482fc82cfbe28e665083795ae
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index ffd6b52..1344144 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -225,6 +225,7 @@
         "androidx.exifinterface_exifinterface",
         "kotlinx-coroutines-android",
         "kotlinx-coroutines-core",
+        "kotlinx_coroutines_test",
         "iconloader_base",
         "SystemUI-tags",
         "SystemUI-proto",
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt
new file mode 100644
index 0000000..e364918
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/RepeatWhenAttached.kt
@@ -0,0 +1,183 @@
+/*
+ *  Copyright (C) 2022 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.lifecycle
+
+import android.view.View
+import android.view.ViewTreeObserver
+import androidx.annotation.MainThread
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.LifecycleRegistry
+import androidx.lifecycle.lifecycleScope
+import com.android.systemui.util.Assert
+import kotlin.coroutines.CoroutineContext
+import kotlin.coroutines.EmptyCoroutineContext
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.DisposableHandle
+import kotlinx.coroutines.launch
+
+/**
+ * Runs the given [block] every time the [View] becomes attached (or immediately after calling this
+ * function, if the view was already attached), automatically canceling the work when the `View`
+ * becomes detached.
+ *
+ * Only use from the main thread.
+ *
+ * When [block] is run, it is run in the context of a [ViewLifecycleOwner] which the caller can use
+ * to launch jobs, with confidence that the jobs will be properly canceled when the view is
+ * detached.
+ *
+ * The [block] may be run multiple times, running once per every time the view is attached. Each
+ * time the block is run for a new attachment event, the [ViewLifecycleOwner] provided will be a
+ * fresh one.
+ *
+ * @param coroutineContext An optional [CoroutineContext] to replace the dispatcher [block] is
+ * invoked on.
+ * @param block The block of code that should be run when the view becomes attached. It can end up
+ * being invoked multiple times if the view is reattached after being detached.
+ * @return A [DisposableHandle] to invoke when the caller of the function destroys its [View] and is
+ * no longer interested in the [block] being run the next time its attached. Calling this is an
+ * optional optimization as the logic will be properly cleaned up and destroyed each time the view
+ * is detached. Using this is not *thread-safe* and should only be used on the main thread.
+ */
+@MainThread
+fun View.repeatWhenAttached(
+    coroutineContext: CoroutineContext = EmptyCoroutineContext,
+    block: suspend LifecycleOwner.(View) -> Unit,
+): DisposableHandle {
+    Assert.isMainThread()
+    val view = this
+    // The suspend block will run on the app's main thread unless the caller supplies a different
+    // dispatcher to use. We don't want it to run on the Dispatchers.Default thread pool as
+    // default behavior. Instead, we want it to run on the view's UI thread since the user will
+    // presumably want to call view methods that require being called from said UI thread.
+    val lifecycleCoroutineContext = Dispatchers.Main + coroutineContext
+    var lifecycleOwner: ViewLifecycleOwner? = null
+    val onAttachListener =
+        object : View.OnAttachStateChangeListener {
+            override fun onViewAttachedToWindow(v: View?) {
+                Assert.isMainThread()
+                lifecycleOwner?.onDestroy()
+                lifecycleOwner =
+                    createLifecycleOwnerAndRun(
+                        view,
+                        lifecycleCoroutineContext,
+                        block,
+                    )
+            }
+
+            override fun onViewDetachedFromWindow(v: View?) {
+                lifecycleOwner?.onDestroy()
+                lifecycleOwner = null
+            }
+        }
+
+    addOnAttachStateChangeListener(onAttachListener)
+    if (view.isAttachedToWindow) {
+        lifecycleOwner =
+            createLifecycleOwnerAndRun(
+                view,
+                lifecycleCoroutineContext,
+                block,
+            )
+    }
+
+    return object : DisposableHandle {
+        override fun dispose() {
+            Assert.isMainThread()
+
+            lifecycleOwner?.onDestroy()
+            lifecycleOwner = null
+            view.removeOnAttachStateChangeListener(onAttachListener)
+        }
+    }
+}
+
+private fun createLifecycleOwnerAndRun(
+    view: View,
+    coroutineContext: CoroutineContext,
+    block: suspend LifecycleOwner.(View) -> Unit,
+): ViewLifecycleOwner {
+    return ViewLifecycleOwner(view).apply {
+        onCreate()
+        lifecycleScope.launch(coroutineContext) { block(view) }
+    }
+}
+
+/**
+ * A [LifecycleOwner] for a [View] for exclusive use by the [repeatWhenAttached] extension function.
+ *
+ * The implementation requires the caller to call [onCreate] and [onDestroy] when the view is
+ * attached to or detached from a view hierarchy. After [onCreate] and before [onDestroy] is called,
+ * the implementation monitors window state in the following way
+ *
+ * * If the window is not visible, we are in the [Lifecycle.State.CREATED] state
+ * * If the window is visible but not focused, we are in the [Lifecycle.State.STARTED] state
+ * * If the window is visible and focused, we are in the [Lifecycle.State.RESUMED] state
+ *
+ * Or in table format:
+ * ```
+ * ┌───────────────┬───────────────────┬──────────────┬─────────────────┐
+ * │ View attached │ Window Visibility │ Window Focus │ Lifecycle State │
+ * ├───────────────┼───────────────────┴──────────────┼─────────────────┤
+ * │ Not attached  │                 Any              │       N/A       │
+ * ├───────────────┼───────────────────┬──────────────┼─────────────────┤
+ * │               │    Not visible    │     Any      │     CREATED     │
+ * │               ├───────────────────┼──────────────┼─────────────────┤
+ * │   Attached    │                   │   No focus   │     STARTED     │
+ * │               │      Visible      ├──────────────┼─────────────────┤
+ * │               │                   │  Has focus   │     RESUMED     │
+ * └───────────────┴───────────────────┴──────────────┴─────────────────┘
+ * ```
+ */
+private class ViewLifecycleOwner(
+    private val view: View,
+) : LifecycleOwner {
+
+    private val windowVisibleListener =
+        ViewTreeObserver.OnWindowVisibilityChangeListener { updateState() }
+    private val windowFocusListener = ViewTreeObserver.OnWindowFocusChangeListener { updateState() }
+
+    private val registry = LifecycleRegistry(this)
+
+    fun onCreate() {
+        registry.currentState = Lifecycle.State.CREATED
+        view.viewTreeObserver.addOnWindowVisibilityChangeListener(windowVisibleListener)
+        view.viewTreeObserver.addOnWindowFocusChangeListener(windowFocusListener)
+        updateState()
+    }
+
+    fun onDestroy() {
+        view.viewTreeObserver.removeOnWindowVisibilityChangeListener(windowVisibleListener)
+        view.viewTreeObserver.removeOnWindowFocusChangeListener(windowFocusListener)
+        registry.currentState = Lifecycle.State.DESTROYED
+    }
+
+    override fun getLifecycle(): Lifecycle {
+        return registry
+    }
+
+    private fun updateState() {
+        registry.currentState =
+            when {
+                view.windowVisibility != View.VISIBLE -> Lifecycle.State.CREATED
+                !view.hasWindowFocus() -> Lifecycle.State.STARTED
+                else -> Lifecycle.State.RESUMED
+            }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwner.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwner.kt
deleted file mode 100644
index 55c7ac9..0000000
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwner.kt
+++ /dev/null
@@ -1,114 +0,0 @@
-package com.android.systemui.lifecycle
-
-import android.view.View
-import android.view.ViewTreeObserver
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.LifecycleRegistry
-
-/**
- * [LifecycleOwner] for Window-added Views.
- *
- * These are [View] instances that are added to a `Window` using the `WindowManager` API.
- *
- * This implementation goes to:
- * * The <b>CREATED</b> `Lifecycle.State` when the view gets attached to the window but the window
- * is not yet visible
- * * The <b>STARTED</b> `Lifecycle.State` when the view is attached to the window and the window is
- * visible
- * * The <b>RESUMED</b> `Lifecycle.State` when the view is attached to the window and the window is
- * visible and the window receives focus
- *
- * In table format:
- * ```
- * | ----------------------------------------------------------------------------- |
- * | View attached to window | Window visible | Window has focus | Lifecycle state |
- * | ----------------------------------------------------------------------------- |
- * |       not attached      |               Any                 |   INITIALIZED   |
- * | ----------------------------------------------------------------------------- |
- * |                         |  not visible   |       Any        |     CREATED     |
- * |                         ----------------------------------------------------- |
- * |        attached         |                |    not focused   |     STARTED     |
- * |                         |   is visible   |----------------------------------- |
- * |                         |                |    has focus     |     RESUMED     |
- * | ----------------------------------------------------------------------------- |
- * ```
- * ### Notes
- * * [dispose] must be invoked when the [LifecycleOwner] is done and won't be reused
- * * It is always better for [LifecycleOwner] implementations to be more explicit than just
- * listening to the state of the `Window`. E.g. if the code that added the `View` to the `Window`
- * already has access to the correct state to know when that `View` should become visible and when
- * it is ready to receive interaction from the user then it already knows when to move to `STARTED`
- * and `RESUMED`, respectively. In that case, it's better to implement your own `LifecycleOwner`
- * instead of relying on the `Window` callbacks.
- */
-class WindowAddedViewLifecycleOwner
-@JvmOverloads
-constructor(
-    private val view: View,
-    registryFactory: (LifecycleOwner) -> LifecycleRegistry = { LifecycleRegistry(it) },
-) : LifecycleOwner {
-
-    private val windowAttachListener =
-        object : ViewTreeObserver.OnWindowAttachListener {
-            override fun onWindowAttached() {
-                updateCurrentState()
-            }
-
-            override fun onWindowDetached() {
-                updateCurrentState()
-            }
-        }
-    private val windowFocusListener =
-        ViewTreeObserver.OnWindowFocusChangeListener { updateCurrentState() }
-    private val windowVisibilityListener =
-        ViewTreeObserver.OnWindowVisibilityChangeListener { updateCurrentState() }
-
-    private val registry = registryFactory(this)
-
-    init {
-        setCurrentState(Lifecycle.State.INITIALIZED)
-
-        with(view.viewTreeObserver) {
-            addOnWindowAttachListener(windowAttachListener)
-            addOnWindowVisibilityChangeListener(windowVisibilityListener)
-            addOnWindowFocusChangeListener(windowFocusListener)
-        }
-
-        updateCurrentState()
-    }
-
-    override fun getLifecycle(): Lifecycle {
-        return registry
-    }
-
-    /**
-     * Disposes of this [LifecycleOwner], performing proper clean-up.
-     *
-     * <p>Invoke this when the instance is finished and won't be reused.
-     */
-    fun dispose() {
-        with(view.viewTreeObserver) {
-            removeOnWindowAttachListener(windowAttachListener)
-            removeOnWindowVisibilityChangeListener(windowVisibilityListener)
-            removeOnWindowFocusChangeListener(windowFocusListener)
-        }
-    }
-
-    private fun updateCurrentState() {
-        val state =
-            when {
-                !view.isAttachedToWindow -> Lifecycle.State.INITIALIZED
-                view.windowVisibility != View.VISIBLE -> Lifecycle.State.CREATED
-                !view.hasWindowFocus() -> Lifecycle.State.STARTED
-                else -> Lifecycle.State.RESUMED
-            }
-        setCurrentState(state)
-    }
-
-    private fun setCurrentState(state: Lifecycle.State) {
-        if (registry.currentState != state) {
-            registry.currentState = state
-        }
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
index 4b71b2c..15e1129 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
@@ -44,8 +44,6 @@
 import android.view.WindowManager.LayoutParams;
 import android.view.WindowManagerGlobal;
 
-import androidx.lifecycle.ViewTreeLifecycleOwner;
-
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.Dumpable;
 import com.android.systemui.R;
@@ -54,7 +52,6 @@
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.KeyguardViewMediator;
-import com.android.systemui.lifecycle.WindowAddedViewLifecycleOwner;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.plugins.statusbar.StatusBarStateController.StateListener;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
@@ -251,15 +248,6 @@
 
         mWindowManager.addView(mNotificationShadeView, mLp);
 
-        // Set up and "inject" a LifecycleOwner bound to the Window-View relationship such that all
-        // views in the sub-tree rooted under this view can access the LifecycleOwner using
-        // ViewTreeLifecycleOwner.get(...).
-        if (ViewTreeLifecycleOwner.get(mNotificationShadeView) == null) {
-            ViewTreeLifecycleOwner.set(
-                    mNotificationShadeView,
-                    new WindowAddedViewLifecycleOwner(mNotificationShadeView));
-        }
-
         mLpChanged.copyFrom(mLp);
         onThemeChanged();
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/RepeatWhenAttachedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/RepeatWhenAttachedTest.kt
new file mode 100644
index 0000000..80f3e46
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/RepeatWhenAttachedTest.kt
@@ -0,0 +1,319 @@
+/*
+ *  Copyright (C) 2022 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.lifecycle
+
+import android.testing.TestableLooper.RunWithLooper
+import android.view.View
+import android.view.ViewTreeObserver
+import androidx.arch.core.executor.ArchTaskExecutor
+import androidx.arch.core.executor.TaskExecutor
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.Assert
+import com.android.systemui.util.mockito.argumentCaptor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.DisposableHandle
+import kotlinx.coroutines.test.runBlockingTest
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TestWatcher
+import org.junit.runner.Description
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mock
+import org.mockito.Mockito.any
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.`when` as whenever
+import org.mockito.junit.MockitoJUnit
+
+@SmallTest
+@RunWith(JUnit4::class)
+@RunWithLooper
+class RepeatWhenAttachedTest : SysuiTestCase() {
+
+    @JvmField @Rule val mockito = MockitoJUnit.rule()
+    @JvmField @Rule val instantTaskExecutor = InstantTaskExecutorRule()
+
+    @Mock private lateinit var view: View
+    @Mock private lateinit var viewTreeObserver: ViewTreeObserver
+
+    private lateinit var block: Block
+    private lateinit var attachListeners: MutableList<View.OnAttachStateChangeListener>
+
+    @Before
+    fun setUp() {
+        Assert.setTestThread(Thread.currentThread())
+        whenever(view.viewTreeObserver).thenReturn(viewTreeObserver)
+        whenever(view.windowVisibility).thenReturn(View.GONE)
+        whenever(view.hasWindowFocus()).thenReturn(false)
+        attachListeners = mutableListOf()
+        whenever(view.addOnAttachStateChangeListener(any())).then {
+            attachListeners.add(it.arguments[0] as View.OnAttachStateChangeListener)
+        }
+        whenever(view.removeOnAttachStateChangeListener(any())).then {
+            attachListeners.remove(it.arguments[0] as View.OnAttachStateChangeListener)
+        }
+        block = Block()
+    }
+
+    @Test(expected = IllegalStateException::class)
+    fun `repeatWhenAttached - enforces main thread`() = runBlockingTest {
+        Assert.setTestThread(null)
+
+        repeatWhenAttached()
+    }
+
+    @Test(expected = IllegalStateException::class)
+    fun `repeatWhenAttached - dispose enforces main thread`() = runBlockingTest {
+        val disposableHandle = repeatWhenAttached()
+        Assert.setTestThread(null)
+
+        disposableHandle.dispose()
+    }
+
+    @Test
+    fun `repeatWhenAttached - view starts detached - runs block when attached`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(false)
+        repeatWhenAttached()
+        assertThat(block.invocationCount).isEqualTo(0)
+
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        attachListeners.last().onViewAttachedToWindow(view)
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.CREATED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - view already attached - immediately runs block`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+
+        repeatWhenAttached()
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.CREATED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - starts visible without focus - STARTED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
+
+        repeatWhenAttached()
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.STARTED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - starts with focus but invisible - CREATED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        whenever(view.hasWindowFocus()).thenReturn(true)
+
+        repeatWhenAttached()
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.CREATED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - starts visible and with focus - RESUMED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
+        whenever(view.hasWindowFocus()).thenReturn(true)
+
+        repeatWhenAttached()
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.RESUMED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - becomes visible without focus - STARTED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        repeatWhenAttached()
+        val listenerCaptor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
+        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(listenerCaptor.capture())
+
+        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
+        listenerCaptor.value.onWindowVisibilityChanged(View.VISIBLE)
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.STARTED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - gains focus but invisible - CREATED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        repeatWhenAttached()
+        val listenerCaptor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
+        verify(viewTreeObserver).addOnWindowFocusChangeListener(listenerCaptor.capture())
+
+        whenever(view.hasWindowFocus()).thenReturn(true)
+        listenerCaptor.value.onWindowFocusChanged(true)
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.CREATED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - becomes visible and gains focus - RESUMED`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        repeatWhenAttached()
+        val visibleCaptor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
+        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(visibleCaptor.capture())
+        val focusCaptor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
+        verify(viewTreeObserver).addOnWindowFocusChangeListener(focusCaptor.capture())
+
+        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
+        visibleCaptor.value.onWindowVisibilityChanged(View.VISIBLE)
+        whenever(view.hasWindowFocus()).thenReturn(true)
+        focusCaptor.value.onWindowFocusChanged(true)
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.RESUMED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - view gets detached - destroys the lifecycle`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        repeatWhenAttached()
+
+        whenever(view.isAttachedToWindow).thenReturn(false)
+        attachListeners.last().onViewDetachedFromWindow(view)
+
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.DESTROYED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - view gets reattached - recreates a lifecycle`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        repeatWhenAttached()
+        whenever(view.isAttachedToWindow).thenReturn(false)
+        attachListeners.last().onViewDetachedFromWindow(view)
+
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        attachListeners.last().onViewAttachedToWindow(view)
+
+        assertThat(block.invocationCount).isEqualTo(2)
+        assertThat(block.invocations[0].lifecycleState).isEqualTo(Lifecycle.State.DESTROYED)
+        assertThat(block.invocations[1].lifecycleState).isEqualTo(Lifecycle.State.CREATED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - dispose attached`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        val handle = repeatWhenAttached()
+
+        handle.dispose()
+
+        assertThat(attachListeners).isEmpty()
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.DESTROYED)
+    }
+
+    @Test
+    fun `repeatWhenAttached - dispose never attached`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(false)
+        val handle = repeatWhenAttached()
+
+        handle.dispose()
+
+        assertThat(attachListeners).isEmpty()
+        assertThat(block.invocationCount).isEqualTo(0)
+    }
+
+    @Test
+    fun `repeatWhenAttached - dispose previously attached now detached`() = runBlockingTest {
+        whenever(view.isAttachedToWindow).thenReturn(true)
+        val handle = repeatWhenAttached()
+        attachListeners.last().onViewDetachedFromWindow(view)
+
+        handle.dispose()
+
+        assertThat(attachListeners).isEmpty()
+        assertThat(block.invocationCount).isEqualTo(1)
+        assertThat(block.latestLifecycleState).isEqualTo(Lifecycle.State.DESTROYED)
+    }
+
+    private fun CoroutineScope.repeatWhenAttached(): DisposableHandle {
+        return view.repeatWhenAttached(
+            coroutineContext = coroutineContext,
+            block = block,
+        )
+    }
+
+    private class Block : suspend LifecycleOwner.(View) -> Unit {
+        data class Invocation(
+            val lifecycleOwner: LifecycleOwner,
+        ) {
+            val lifecycleState: Lifecycle.State
+                get() = lifecycleOwner.lifecycle.currentState
+        }
+
+        private val _invocations = mutableListOf<Invocation>()
+        val invocations: List<Invocation> = _invocations
+        val invocationCount: Int
+            get() = _invocations.size
+        val latestLifecycleState: Lifecycle.State
+            get() = _invocations.last().lifecycleState
+
+        override suspend fun invoke(lifecycleOwner: LifecycleOwner, view: View) {
+            _invocations.add(Invocation(lifecycleOwner))
+        }
+    }
+
+    /**
+     * Test rule that makes ArchTaskExecutor main thread assertions pass. There is one such assert
+     * in LifecycleRegistry.
+     */
+    class InstantTaskExecutorRule : TestWatcher() {
+        // TODO(b/240620122): This is a copy of
+        //  androidx/arch/core/executor/testing/InstantTaskExecutorRule which should be replaced
+        //  with a dependency on the real library once b/ is cleared.
+        override fun starting(description: Description) {
+            super.starting(description)
+            ArchTaskExecutor.getInstance()
+                .setDelegate(
+                    object : TaskExecutor() {
+                        override fun executeOnDiskIO(runnable: Runnable) {
+                            runnable.run()
+                        }
+
+                        override fun postToMainThread(runnable: Runnable) {
+                            runnable.run()
+                        }
+
+                        override fun isMainThread(): Boolean {
+                            return true
+                        }
+                    }
+                )
+        }
+
+        override fun finished(description: Description) {
+            super.finished(description)
+            ArchTaskExecutor.getInstance().setDelegate(null)
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwnerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwnerTest.kt
deleted file mode 100644
index 4f5c570..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/WindowAddedViewLifecycleOwnerTest.kt
+++ /dev/null
@@ -1,150 +0,0 @@
-package com.android.systemui.lifecycle
-
-import android.view.View
-import android.view.ViewTreeObserver
-import androidx.lifecycle.Lifecycle
-import androidx.lifecycle.LifecycleRegistry
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.argumentCaptor
-import com.android.systemui.util.mockito.capture
-import com.google.common.truth.Truth.assertThat
-import org.junit.Before
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
-
-@SmallTest
-@RunWith(JUnit4::class)
-class WindowAddedViewLifecycleOwnerTest : SysuiTestCase() {
-
-    @Mock lateinit var view: View
-    @Mock lateinit var viewTreeObserver: ViewTreeObserver
-
-    private lateinit var underTest: WindowAddedViewLifecycleOwner
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-        whenever(view.viewTreeObserver).thenReturn(viewTreeObserver)
-        whenever(view.isAttachedToWindow).thenReturn(false)
-        whenever(view.windowVisibility).thenReturn(View.INVISIBLE)
-        whenever(view.hasWindowFocus()).thenReturn(false)
-
-        underTest = WindowAddedViewLifecycleOwner(view) { LifecycleRegistry.createUnsafe(it) }
-    }
-
-    @Test
-    fun `detached - invisible - does not have focus -- INITIALIZED`() {
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.INITIALIZED)
-    }
-
-    @Test
-    fun `detached - invisible - has focus -- INITIALIZED`() {
-        whenever(view.hasWindowFocus()).thenReturn(true)
-        val captor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
-        verify(viewTreeObserver).addOnWindowFocusChangeListener(capture(captor))
-        captor.value.onWindowFocusChanged(true)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.INITIALIZED)
-    }
-
-    @Test
-    fun `detached - visible - does not have focus -- INITIALIZED`() {
-        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
-        val captor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
-        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(capture(captor))
-        captor.value.onWindowVisibilityChanged(View.VISIBLE)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.INITIALIZED)
-    }
-
-    @Test
-    fun `detached - visible - has focus -- INITIALIZED`() {
-        whenever(view.hasWindowFocus()).thenReturn(true)
-        val focusCaptor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
-        verify(viewTreeObserver).addOnWindowFocusChangeListener(capture(focusCaptor))
-        focusCaptor.value.onWindowFocusChanged(true)
-
-        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
-        val visibilityCaptor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
-        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(capture(visibilityCaptor))
-        visibilityCaptor.value.onWindowVisibilityChanged(View.VISIBLE)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.INITIALIZED)
-    }
-
-    @Test
-    fun `attached - invisible - does not have focus -- CREATED`() {
-        whenever(view.isAttachedToWindow).thenReturn(true)
-        val captor = argumentCaptor<ViewTreeObserver.OnWindowAttachListener>()
-        verify(viewTreeObserver).addOnWindowAttachListener(capture(captor))
-        captor.value.onWindowAttached()
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.CREATED)
-    }
-
-    @Test
-    fun `attached - invisible - has focus -- CREATED`() {
-        whenever(view.isAttachedToWindow).thenReturn(true)
-        val attachCaptor = argumentCaptor<ViewTreeObserver.OnWindowAttachListener>()
-        verify(viewTreeObserver).addOnWindowAttachListener(capture(attachCaptor))
-        attachCaptor.value.onWindowAttached()
-
-        whenever(view.hasWindowFocus()).thenReturn(true)
-        val focusCaptor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
-        verify(viewTreeObserver).addOnWindowFocusChangeListener(capture(focusCaptor))
-        focusCaptor.value.onWindowFocusChanged(true)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.CREATED)
-    }
-
-    @Test
-    fun `attached - visible - does not have focus -- STARTED`() {
-        whenever(view.isAttachedToWindow).thenReturn(true)
-        val attachCaptor = argumentCaptor<ViewTreeObserver.OnWindowAttachListener>()
-        verify(viewTreeObserver).addOnWindowAttachListener(capture(attachCaptor))
-        attachCaptor.value.onWindowAttached()
-
-        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
-        val visibilityCaptor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
-        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(capture(visibilityCaptor))
-        visibilityCaptor.value.onWindowVisibilityChanged(View.VISIBLE)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.STARTED)
-    }
-
-    @Test
-    fun `attached - visible - has focus -- RESUMED`() {
-        whenever(view.isAttachedToWindow).thenReturn(true)
-        val attachCaptor = argumentCaptor<ViewTreeObserver.OnWindowAttachListener>()
-        verify(viewTreeObserver).addOnWindowAttachListener(capture(attachCaptor))
-        attachCaptor.value.onWindowAttached()
-
-        whenever(view.hasWindowFocus()).thenReturn(true)
-        val focusCaptor = argumentCaptor<ViewTreeObserver.OnWindowFocusChangeListener>()
-        verify(viewTreeObserver).addOnWindowFocusChangeListener(capture(focusCaptor))
-        focusCaptor.value.onWindowFocusChanged(true)
-
-        whenever(view.windowVisibility).thenReturn(View.VISIBLE)
-        val visibilityCaptor = argumentCaptor<ViewTreeObserver.OnWindowVisibilityChangeListener>()
-        verify(viewTreeObserver).addOnWindowVisibilityChangeListener(capture(visibilityCaptor))
-        visibilityCaptor.value.onWindowVisibilityChanged(View.VISIBLE)
-
-        assertThat(underTest.lifecycle.currentState).isEqualTo(Lifecycle.State.RESUMED)
-    }
-
-    @Test
-    fun dispose() {
-        underTest.dispose()
-
-        verify(viewTreeObserver).removeOnWindowAttachListener(any())
-        verify(viewTreeObserver).removeOnWindowVisibilityChangeListener(any())
-        verify(viewTreeObserver).removeOnWindowFocusChangeListener(any())
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
index ec1fa48..ad3d3d2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
@@ -28,7 +28,6 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.clearInvocations;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
@@ -44,8 +43,6 @@
 import android.view.View;
 import android.view.WindowManager;
 
-import androidx.lifecycle.LifecycleOwner;
-import androidx.lifecycle.ViewTreeLifecycleOwner;
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.colorextraction.ColorExtractor;
@@ -188,24 +185,6 @@
     }
 
     @Test
-    public void attach_setsUpLifecycleOwner() {
-        mNotificationShadeWindowController.attach();
-
-        assertThat(ViewTreeLifecycleOwner.get(mNotificationShadeWindowView)).isNotNull();
-    }
-
-    @Test
-    public void attach_doesNotSetUpLifecycleOwnerIfAlreadySet() {
-        final LifecycleOwner previouslySet = mock(LifecycleOwner.class);
-        ViewTreeLifecycleOwner.set(mNotificationShadeWindowView, previouslySet);
-
-        mNotificationShadeWindowController.attach();
-
-        assertThat(ViewTreeLifecycleOwner.get(mNotificationShadeWindowView))
-                .isEqualTo(previouslySet);
-    }
-
-    @Test
     public void setScrimsVisibility_earlyReturn() {
         clearInvocations(mWindowManager);
         mNotificationShadeWindowController.setScrimsVisibility(ScrimController.TRANSPARENT);