Only capture test snapshots after a frame completes

redwood-testing/src/commonMain/kotlin/app/cash/redwood/testing/TestRedwoodComposition.kt

@@ -29,12 +29,9 @@ import app.cash.redwood.widget.WidgetSystem
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.TimeoutCancellationException
-import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.plus
 import kotlinx.coroutines.withTimeout
 
@@ -94,16 +91,14 @@ private class RealTestRedwoodComposition<W : Any, S>(
   onBackPressedDispatcher: OnBackPressedDispatcher,
   savedState: TestSavedState?,
   initialUiConfiguration: UiConfiguration,
-  createSnapshot: () -> S,
+  private val createSnapshot: () -> S,
 ) : TestRedwoodComposition<S> {
-  /** Emit frames manually in [sendFrames]. */
+  /** Emits frames manually in [awaitSnapshot]. */
   private val clock = BroadcastFrameClock()
   private var timeNanos = 0L
   private val frameDelay = 1.seconds / 60
   private var contentSet = false
-
-  /** Channel with the most recent snapshot, if any. */
-  private val snapshots = Channel<S>(Channel.CONFLATED)
+  private var hasChanges = false
 
   override val uiConfigurations = MutableStateFlow(initialUiConfiguration)
 
@@ -125,12 +120,7 @@ private class RealTestRedwoodComposition<W : Any, S>(
     saveableStateRegistry = savedStateRegistry,
     uiConfigurations = uiConfigurations,
     widgetSystem = widgetSystem,
-    onEndChanges = {
-      val newSnapshot = createSnapshot()
-
-      // trySend always succeeds on a CONFLATED channel.
-      check(snapshots.trySend(newSnapshot).isSuccess)
-    },
+    onEndChanges = { hasChanges = true },
   )
 
   override fun setContent(content: @Composable () -> Unit) {
@@ -141,26 +131,19 @@ private class RealTestRedwoodComposition<W : Any, S>(
   override suspend fun awaitSnapshot(timeout: Duration): S {
     check(contentSet) { "setContent must be called first!" }
 
-    // Await at least one change, sending frames while we wait.
-    return withTimeout(timeout) {
-      val sendFramesJob = sendFrames()
-      try {
-        snapshots.receive()
-      } finally {
-        sendFramesJob.cancel()
-      }
-    }
-  }
-
-  /** Launches a job that sends a frame immediately and again every 16 ms until it's canceled. */
-  private fun CoroutineScope.sendFrames(): Job {
-    return launch {
+    // Await changes, sending at least one frame while we wait.
+    withTimeout(timeout) {
       while (true) {
         clock.sendFrame(timeNanos)
+        if (hasChanges) break
+
         timeNanos += frameDelay.inWholeNanoseconds
         delay(frameDelay)
       }
     }
+
+    hasChanges = false
+    return createSnapshot()
   }
 
   override fun cancel() {

redwood-testing/src/commonTest/kotlin/app/cash/redwood/testing/TestRedwoodCompositionTest.kt

@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2024 Square, Inc.
+ *
+ * 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 app.cash.redwood.testing
+
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.movableContentOf
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import app.cash.redwood.layout.compose.Column
+import app.cash.redwood.layout.compose.Row
+import app.cash.redwood.layout.testing.RedwoodLayoutTestingWidgetFactory
+import app.cash.redwood.lazylayout.testing.RedwoodLazyLayoutTestingWidgetFactory
+import app.cash.redwood.widget.MutableListChildren
+import assertk.assertThat
+import assertk.assertions.isEqualTo
+import com.example.redwood.testapp.compose.Text
+import com.example.redwood.testapp.testing.TestSchemaTestingWidgetFactory
+import com.example.redwood.testapp.widget.TestSchemaWidgetSystem
+import kotlin.test.Test
+import kotlinx.coroutines.test.runTest
+
+class TestRedwoodCompositionTest {
+  @Test fun awaitSnapshotCapturesMultipleChanges() = runTest {
+    var count = 0
+    val tester = TestRedwoodComposition(
+      scope = backgroundScope,
+      widgetSystem = TestSchemaWidgetSystem(
+        TestSchema = TestSchemaTestingWidgetFactory(),
+        RedwoodLayout = RedwoodLayoutTestingWidgetFactory(),
+        RedwoodLazyLayout = RedwoodLazyLayoutTestingWidgetFactory(),
+      ),
+      container = MutableListChildren<WidgetValue>(),
+      createSnapshot = { ++count },
+    )
+
+    // The content of a movableContentOf is applied to the node tree separately, resulting in
+    // two calls to Applier.onEndChanges. If this signal is used to emit a snapshot, only a
+    // partial view of the recomposition will be available.
+    var isRow by mutableStateOf(true)
+    tester.setContent {
+      val movable = remember {
+        movableContentOf {
+          Text("one")
+        }
+      }
+      if (isRow) {
+        Row { movable() }
+      } else {
+        Column { movable() }
+      }
+    }
+
+    assertThat(tester.awaitSnapshot()).isEqualTo(1)
+    isRow = false
+    assertThat(tester.awaitSnapshot()).isEqualTo(2)
+  }
+}