Test that UIs are closed properly (#1644)

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/CodeHost.kt

@@ -15,16 +15,14 @@
  */
 package app.cash.redwood.treehouse
 
-import app.cash.zipline.ZiplineScope
-
 /** Manages loading and hot-reloading a series of code sessions. */
 internal interface CodeHost<A : AppService> {
   val stateStore: StateStore
 
   /** Only accessed on [TreehouseDispatchers.ui]. */
   val session: CodeSession<A>?
 
-  fun applyZiplineScope(appService: A, ziplineScope: ZiplineScope): A
+  fun newServiceScope(): ServiceScope<A>
 
   fun addListener(listener: Listener<A>)
 
@@ -33,4 +31,17 @@ internal interface CodeHost<A : AppService> {
   interface Listener<A : AppService> {
     fun codeSessionChanged(next: CodeSession<A>)
   }
+
+  /**
+   * Tracks all of the services created to produce a UI, and offers a single mechanism to close
+   * them all. Note that closing this does not close the app services it was applied to.
+   */
+  interface ServiceScope<A : AppService> {
+    /**
+     * Returns a new instance that forwards calls to [appService] and keeps track of returned
+     * instances so they may be closed.
+     */
+    fun apply(appService: A): A
+    fun close()
+  }
 }

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseApp.kt

@@ -199,8 +199,18 @@ public class TreehouseApp<A : AppService> private constructor(
      */
     private val listeners = mutableListOf<CodeHost.Listener<A>>()
 
-    override fun applyZiplineScope(appService: A, ziplineScope: ZiplineScope): A {
-      return appService.withScope(ziplineScope)
+    override fun newServiceScope(): CodeHost.ServiceScope<A> {
+      val ziplineScope = ZiplineScope()
+
+      return object : CodeHost.ServiceScope<A> {
+        override fun apply(appService: A): A {
+          return appService.withScope(ziplineScope)
+        }
+
+        override fun close() {
+          ziplineScope.close()
+        }
+      }
     }
 
     override fun addListener(listener: CodeHost.Listener<A>) {

redwood-treehouse-host/src/commonMain/kotlin/app/cash/redwood/treehouse/TreehouseAppContent.kt

@@ -286,7 +286,7 @@ private class ViewContentCodeBinding<A : AppService>(
   private var bridgeOrNull: ProtocolBridge<*>? = null
 
   /** Only accessed on [TreehouseDispatchers.zipline]. */
-  private val ziplineScope = ZiplineScope()
+  private val serviceScope = codeHost.newServiceScope()
 
   /** Only accessed on [TreehouseDispatchers.zipline]. Null after [cancel]. */
   private var treehouseUiOrNull: ZiplineTreehouseUi? = null
@@ -385,7 +385,7 @@ private class ViewContentCodeBinding<A : AppService>(
 
   fun start(session: CodeSession<A>) {
     bindingScope.launch(dispatchers.zipline) {
-      val scopedAppService = codeHost.applyZiplineScope(session.appService, ziplineScope)
+      val scopedAppService = serviceScope.apply(session.appService)
       val treehouseUi = contentSource.get(scopedAppService)
       treehouseUiOrNull = treehouseUi
       val restoredId = viewOrNull?.stateSnapshotId
@@ -452,7 +452,7 @@ private class ViewContentCodeBinding<A : AppService>(
     appScope.launch(dispatchers.zipline) {
       treehouseUiOrNull = null
       bindingScope.cancel()
-      ziplineScope.close()
+      serviceScope.close()
     }
   }
 }

redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/FakeAppService.kt

@@ -15,17 +15,15 @@
  */
 package app.cash.redwood.treehouse
 
-internal class FakeAppService(
+internal class FakeAppService private constructor(
   private val name: String,
   private val eventLog: EventLog,
+  private val listeners: List<Listener>,
+  private val mutableUis: MutableList<FakeZiplineTreehouseUi>,
 ) : AppService {
-  val uis = mutableListOf<FakeZiplineTreehouseUi>()
 
-  fun newUi(): ZiplineTreehouseUi {
-    val result = FakeZiplineTreehouseUi("$name.uis[${uis.size}]", eventLog)
-    uis += result
-    return result
-  }
+  val uis: List<FakeZiplineTreehouseUi>
+    get() = mutableUis.toList()
 
   override val appLifecycle = object : AppLifecycle {
     override fun start(host: AppLifecycle.Host) {
@@ -35,4 +33,34 @@ internal class FakeAppService(
     override fun sendFrame(timeNanos: Long) {
     }
   }
+
+  constructor(
+    name: String,
+    eventLog: EventLog,
+  ) : this(name, eventLog, listOf(), mutableListOf())
+
+  /**
+   * Return a FakeAppService that shares all state with this, but that notifies [listeners] of
+   * events triggered through it.
+   *
+   * Note that this does not add [listener] to receive events triggered directly on this.
+   *
+   * This awkward pattern emulates `ZiplineScope`, which also only tracks objects on the scoped
+   * wrapper.
+   */
+  fun withListener(listener: Listener) =
+    FakeAppService(name, eventLog, listeners + listener, mutableUis)
+
+  fun newUi(): ZiplineTreehouseUi {
+    val result = FakeZiplineTreehouseUi("$name.uis[${mutableUis.size}]", eventLog)
+    for (listener in listeners) {
+      listener.onNewUi(result)
+    }
+    mutableUis += result
+    return result
+  }
+
+  interface Listener {
+    fun onNewUi(ui: ZiplineTreehouseUi)
+  }
 }

redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/FakeCodeHost.kt

@@ -15,12 +15,10 @@
  */
 package app.cash.redwood.treehouse
 
-import app.cash.zipline.ZiplineScope
-
-internal class FakeCodeHost<A : AppService> : CodeHost<A> {
+internal class FakeCodeHost : CodeHost<FakeAppService> {
   override val stateStore = MemoryStateStore()
 
-  override var session: CodeSession<A>? = null
+  override var session: CodeSession<FakeAppService>? = null
     set(value) {
       require(value != null)
 
@@ -34,15 +32,33 @@ internal class FakeCodeHost<A : AppService> : CodeHost<A> {
       field = value
     }
 
-  private val listeners = mutableListOf<CodeHost.Listener<A>>()
+  private val listeners = mutableListOf<CodeHost.Listener<FakeAppService>>()
+
+  override fun newServiceScope(): CodeHost.ServiceScope<FakeAppService> {
+    return object : CodeHost.ServiceScope<FakeAppService> {
+      val uisToClose = mutableListOf<ZiplineTreehouseUi>()
+
+      override fun apply(appService: FakeAppService): FakeAppService {
+        return appService.withListener(object : FakeAppService.Listener {
+          override fun onNewUi(ui: ZiplineTreehouseUi) {
+            uisToClose += ui
+          }
+        })
+      }
 
-  override fun applyZiplineScope(appService: A, ziplineScope: ZiplineScope) = appService
+      override fun close() {
+        for (ui in uisToClose) {
+          ui.close()
+        }
+      }
+    }
+  }
 
-  override fun addListener(listener: CodeHost.Listener<A>) {
+  override fun addListener(listener: CodeHost.Listener<FakeAppService>) {
     listeners += listener
   }
 
-  override fun removeListener(listener: CodeHost.Listener<A>) {
+  override fun removeListener(listener: CodeHost.Listener<FakeAppService>) {
     listeners -= listener
   }
 }

redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/FakeZiplineTreehouseUi.kt

@@ -72,4 +72,8 @@ class FakeZiplineTreehouseUi(
   ) {
     error("unexpected call")
   }
+
+  override fun close() {
+    eventLog += "$name.close()"
+  }
 }

redwood-treehouse-host/src/commonTest/kotlin/app/cash/redwood/treehouse/TreehouseAppContentTest.kt

@@ -31,7 +31,7 @@ class TreehouseAppContentTest {
   private val eventLog = EventLog()
 
   private val dispatcher = UnconfinedTestDispatcher()
-  private val codeHost = FakeCodeHost<FakeAppService>()
+  private val codeHost = FakeCodeHost()
   private val dispatchers = FakeDispatchers(dispatcher, dispatcher)
   private val eventPublisher = FakeEventPublisher()
   private val codeListener = FakeCodeListener(eventLog)
@@ -59,6 +59,7 @@ class TreehouseAppContentTest {
     assertThat(view1.children.single().value.label).isEqualTo("hello")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
   }
 
   @Test
@@ -81,6 +82,7 @@ class TreehouseAppContentTest {
     eventLog.takeEvent("codeListener.onCodeLoaded(view1, initial = true)")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
   }
 
   @Test
@@ -102,6 +104,7 @@ class TreehouseAppContentTest {
     assertThat(view1.children.single().value.label).isEqualTo("hello")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
   }
 
   @Test
@@ -120,6 +123,7 @@ class TreehouseAppContentTest {
     assertThat(view1.children.single().value.label).isEqualTo("hello")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
   }
 
   /** This exercises hot reloading. The view sees new code. */
@@ -131,26 +135,30 @@ class TreehouseAppContentTest {
     content.bind(view1)
     eventLog.takeEvent("codeListener.onInitialCodeLoading(view1)")
 
-    codeHost.session = FakeCodeSession("codeSessionA", eventLog)
-    codeHost.session!!.appService.uis.single().addWidget("helloA")
+    val codeSessionA = FakeCodeSession("codeSessionA", eventLog)
+    codeHost.session = codeSessionA
+    codeSessionA.appService.uis.single().addWidget("helloA")
     eventLog.takeEvent("codeSessionA.start()")
     eventLog.takeEvent("codeSessionA.app.uis[0].start()")
     eventLog.takeEvent("codeListener.onCodeLoaded(view1, initial = true)")
 
-    codeHost.session = FakeCodeSession("codeSessionB", eventLog)
+    val codeSessionB = FakeCodeSession("codeSessionB", eventLog)
+    codeHost.session = codeSessionB
     eventLog.takeEvent("codeSessionA.cancel()")
     eventLog.takeEvent("codeSessionB.start()")
     eventLog.takeEvent("codeSessionB.app.uis[0].start()")
 
-    // No onCodeLoaded() and no reset() until the new code's first widget is added!
+    // This still shows UI from codeSessionA. There's no onCodeLoaded() and no reset() until the new
+    // code's first widget is added!
     assertThat(view1.children.single().value.label).isEqualTo("helloA")
-    assertThat(eventLog.assertNoEvents())
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
 
-    codeHost.session!!.appService.uis.single().addWidget("helloB")
+    codeSessionB.appService.uis.single().addWidget("helloB")
     eventLog.takeEvent("codeListener.onCodeLoaded(view1, initial = false)")
     assertThat(view1.children.single().value.label).isEqualTo("helloB")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionB.app.uis[0].close()")
   }
 
   @Test
@@ -204,6 +212,8 @@ class TreehouseAppContentTest {
     assertThat(view1.children.single().value.label).isEqualTo("helloA")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[0].close()")
+
     content.bind(view1)
     eventLog.takeEvent("codeSessionA.app.uis[1].start()")
 
@@ -212,6 +222,7 @@ class TreehouseAppContentTest {
     assertThat(view1.children.single().value.label).isEqualTo("helloB")
 
     content.unbind()
+    eventLog.takeEvent("codeSessionA.app.uis[1].close()")
   }
 
   private fun TestScope.treehouseAppContent(): TreehouseAppContent<FakeAppService> {