Make HTML TestUtils wait functions cancellable

html/test-utils/src/jsMain/kotlin/org/jetbrains/compose/web/testutils/TestUtils.kt

@@ -4,6 +4,7 @@ import androidx.compose.runtime.Composable
 import androidx.compose.runtime.MonotonicFrameClock
 import kotlinx.browser.document
 import kotlinx.browser.window
+import kotlinx.coroutines.suspendCancellableCoroutine
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.MainScope
 import kotlinx.coroutines.promise
@@ -84,21 +85,31 @@ class TestScope : CoroutineScope by MainScope() {
      * Suspends until [element] observes any change to its html.
      */
     suspend fun waitForChanges(element: HTMLElement = root) {
-        suspendCoroutine<Unit> { continuation ->
+        suspendCancellableCoroutine<Unit> { continuation ->
             val observer = MutationObserver { _, observer ->
                 continuation.resume(Unit)
                 observer.disconnect()
             }
             observer.observe(element, MutationObserverOptions)
+
+            continuation.invokeOnCancellation {
+                observer.disconnect()
+            }
         }
     }
 
     /**
      * Suspends until recomposition completes.
      */
     suspend fun waitForRecompositionComplete() {
-        suspendCoroutine<Unit> { continuation ->
+        suspendCancellableCoroutine<Unit> { continuation ->
             waitForRecompositionCompleteContinuation = continuation
+
+            continuation.invokeOnCancellation {
+                if (waitForRecompositionCompleteContinuation === continuation) {
+                    waitForRecompositionCompleteContinuation = null
+                }
+            }
         }
     }
 }

html/test-utils/src/jsTest/kotlin/TestsForTestUtils.kt

@@ -3,6 +3,8 @@ import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.currentRecomposeScope
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.TimeoutCancellationException
+import kotlinx.coroutines.withTimeout
 import org.jetbrains.compose.web.testutils.ComposeWebExperimentalTestsApi
 import org.jetbrains.compose.web.testutils.runTest
 import kotlin.test.Test
@@ -79,4 +81,54 @@ class TestsForTestUtils {
         assertEquals(true, waitForChangesContinued)
         assertEquals("<div>Hello World!</div>", root.outerHTML)
     }
+
+    @Test
+    fun waitForChanges_cancels_with_timeout() = runTest {
+
+        var cancelled = false
+
+        val job = launch {
+            try {
+                withTimeout(1000) {
+                    waitForChanges(root)
+                }
+            } catch (t: TimeoutCancellationException) {
+                cancelled = true
+                throw t
+            }
+        }
+
+        delay(100) // to check that `waitForChanges` is suspended after delay
+        assertEquals(false, cancelled)
+
+        delay(1000) // to check that `waitForChanges` is cancelled after timeout
+        assertEquals(true, cancelled)
+
+        job.join()
+    }
+
+    @Test
+    fun waitForRecompositionComplete_cancels_with_timeout() = runTest {
+
+        var cancelled = false
+
+        val job = launch {
+            try {
+                withTimeout(1000) {
+                    waitForRecompositionComplete()
+                }
+            } catch (t: TimeoutCancellationException) {
+                cancelled = true
+                throw t
+            }
+        }
+
+        delay(100) // to check that `waitForRecompositionComplete` is suspended after delay
+        assertEquals(false, cancelled)
+
+        delay(1000) // to check that `waitForRecompositionComplete` is cancelled after timeout
+        assertEquals(true, cancelled)
+
+        job.join()
+    }
 }