Implement SemanticsNodeInteraction.captureToImage() (#589)

compose/ui/ui-test-junit4/src/desktopTest/kotlin/androidx/compose/ui/test/BasicTestTest.kt

@@ -16,13 +16,18 @@
 
 package androidx.compose.ui.test
 
+import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.size
 import androidx.compose.material.Text
 import androidx.compose.runtime.*
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import androidx.compose.ui.graphics.toArgb
 import androidx.compose.ui.layout.Layout
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.unit.dp
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertTrue
@@ -111,4 +116,25 @@ class BasicTestTest {
         text = "2"
         rule.onNodeWithTag("text").assertTextEquals("2")
     }
+
+    @Test
+    fun testCaptureToImage() {
+        val color = Color.Green
+        rule.setContent {
+            Box(Modifier.testTag("box").size(20.dp).background(color))
+        }
+
+        val screenshot = rule.onNodeWithTag("box").captureToImage()
+
+        assertEquals(20, screenshot.width)
+        assertEquals(20, screenshot.height)
+
+        IntArray(20*20).let { buffer ->
+            screenshot.readPixels(buffer)
+            val expectedPixel = color.toArgb()
+            for (pixel in buffer) {
+                assertEquals(expectedPixel, pixel)
+            }
+        }
+    }
 }

compose/ui/ui-test-junit4/src/skikoMain/kotlin/androidx/compose/ui/test/ComposeUiTest.skikoMain.kt

@@ -316,15 +316,19 @@ class SkikoComposeUiTest(
         return surface.makeImageSnapshot().toComposeImageBitmap()
     }
 
-    fun SemanticsNodeInteraction.captureToImage(): ImageBitmap {
-        waitForIdle()
-        val rect = fetchSemanticsNode().boundsInWindow
+    fun captureToImage(semanticsNode: SemanticsNode): ImageBitmap {
+        val rect = semanticsNode.boundsInWindow
         val iRect = IRect.makeLTRB(rect.left.toInt(), rect.top.toInt(), rect.right.toInt(), rect.bottom.toInt())
         val image = surface.makeImageSnapshot(iRect)
         return image!!.toComposeImageBitmap()
     }
 
-    private inner class DesktopTestOwner : TestOwner {
+    fun SemanticsNodeInteraction.captureToImage(): ImageBitmap {
+        return captureToImage(fetchSemanticsNode())
+    }
+
+    @OptIn(InternalComposeUiApi::class)
+    private inner class DesktopTestOwner : TestOwner, SkikoTestOwner {
         val roots: Set<RootForTest>
             get() = this@SkikoComposeUiTest.scene.roots
 
@@ -346,6 +350,9 @@ class SkikoComposeUiTest(
 
         override val mainClock get() =
             this@SkikoComposeUiTest.mainClock
+
+        override fun captureToImage(semanticsNode: SemanticsNode): ImageBitmap =
+            this@SkikoComposeUiTest.captureToImage(semanticsNode)
     }
 }
 

compose/ui/ui-test/src/skikoMain/kotlin/androidx/compose/ui/test/SkikoImageHelpers.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2023 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 androidx.compose.ui.test
+
+import androidx.compose.ui.InternalComposeUiApi
+import androidx.compose.ui.graphics.ImageBitmap
+
+
+/**
+ * Captures the underlying semantics node's surface into bitmap.
+ */
+@OptIn(InternalTestApi::class, InternalComposeUiApi::class)
+fun SemanticsNodeInteraction.captureToImage(): ImageBitmap {
+    val semanticsNode = fetchSemanticsNode("Failed to capture a node to bitmap.")
+    return (testContext.testOwner as SkikoTestOwner).captureToImage(semanticsNode)
+}

compose/ui/ui-test/src/skikoMain/kotlin/androidx/compose/ui/test/SkikoTestOwner.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2023 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 androidx.compose.ui.test
+
+import androidx.compose.ui.InternalComposeUiApi
+import androidx.compose.ui.graphics.ImageBitmap
+import androidx.compose.ui.semantics.SemanticsNode
+
+/**
+ * Marker interface to be implemented by the Skiko `TestOwner`. The purpose is to allow code in the
+ * `ui-test` module to access functions specific to the Skiko test implementation (e.g. in
+ * `ui-test-junit4`).
+ */
+// Ideally, this would be @VisibleForTesting, but that's an Android annotation we don't want a
+// dependency on
+@InternalComposeUiApi
+interface SkikoTestOwner {
+
+    fun captureToImage(semanticsNode: SemanticsNode): ImageBitmap
+
+}