diff --git a/include-build/roborazzi-gradle-plugin/src/integrationTest/projects/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/CustomPreviewTester.kt b/include-build/roborazzi-gradle-plugin/src/integrationTest/projects/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/CustomPreviewTester.kt
index 3f3e6b76a..e122cb0b2 100644
--- a/include-build/roborazzi-gradle-plugin/src/integrationTest/projects/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/CustomPreviewTester.kt
+++ b/include-build/roborazzi-gradle-plugin/src/integrationTest/projects/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/CustomPreviewTester.kt
@@ -34,6 +34,6 @@ class CustomPreviewTester : ComposePreviewTester<AndroidPreviewInfo> by AndroidC
     composeTestRule.setContent {
       preview()
     }
-    composeTestRule.onRoot().captureRoboImage("${roborazziSystemPropertyOutputDirectory()}/${preview.methodName}.png")
+    composeTestRule.onRoot().captureRoboImage("${roborazziSystemPropertyOutputDirectory()}/${preview.methodName}.${provideRoborazziContext().imageExtension}")
   }
 }
\ No newline at end of file
diff --git a/roborazzi-compose-preview-scanner-support/src/main/java/com/github/takahirom/roborazzi/RoborazziPreviewScannerSupport.kt b/roborazzi-compose-preview-scanner-support/src/main/java/com/github/takahirom/roborazzi/RoborazziPreviewScannerSupport.kt
index 14ed93a80..020bd2845 100644
--- a/roborazzi-compose-preview-scanner-support/src/main/java/com/github/takahirom/roborazzi/RoborazziPreviewScannerSupport.kt
+++ b/roborazzi-compose-preview-scanner-support/src/main/java/com/github/takahirom/roborazzi/RoborazziPreviewScannerSupport.kt
@@ -108,7 +108,7 @@ class AndroidComposePreviewTester : ComposePreviewTester<AndroidPreviewInfo> {
       preview.declaringClass,
       createScreenshotIdFor(preview)
     )
-    val filePath = "$pathPrefix$name.png"
+    val filePath = "$pathPrefix$name.${provideRoborazziContext().imageExtension}"
     preview.captureRoboImage(filePath)
   }
 
diff --git a/roborazzi-compose/src/main/java/com/github/takahirom/roborazzi/RoborazziCompose.kt b/roborazzi-compose/src/main/java/com/github/takahirom/roborazzi/RoborazziCompose.kt
index 7c024de03..a1a67bb67 100644
--- a/roborazzi-compose/src/main/java/com/github/takahirom/roborazzi/RoborazziCompose.kt
+++ b/roborazzi-compose/src/main/java/com/github/takahirom/roborazzi/RoborazziCompose.kt
@@ -9,6 +9,7 @@ import androidx.compose.ui.platform.ComposeView
 import androidx.compose.ui.platform.ViewRootForTest
 import androidx.test.core.app.ActivityScenario
 import androidx.test.core.app.ApplicationProvider
+import androidx.test.espresso.Espresso
 import org.robolectric.Shadows
 import java.io.File
 
@@ -31,15 +32,26 @@ fun captureRoboImage(
 ) {
   if (!roborazziOptions.taskType.isEnabled()) return
   registerRoborazziActivityToRobolectricIfNeeded()
+  // Views needs to be laid out before we can capture them
+  Espresso.onIdle()
+
   val activityScenario = ActivityScenario.launch(RoborazziTransparentActivity::class.java)
   activityScenario.use {
     activityScenario.onActivity { activity ->
       activity.setContent(content = content)
-      val composeView = activity.window.decorView
-        .findViewById<ViewGroup>(android.R.id.content)
-        .getChildAt(0) as ComposeView
-      val viewRootForTest = composeView.getChildAt(0) as ViewRootForTest
-      viewRootForTest.view.captureRoboImage(file, roborazziOptions)
+      val windowRoots = fetchRobolectricWindowRoots()
+      if (windowRoots.size <= 1) {
+        val composeView = activity.window.decorView
+          .findViewById<ViewGroup>(android.R.id.content)
+          .getChildAt(0) as ComposeView
+        val viewRootForTest = composeView.getChildAt(0) as ViewRootForTest
+        viewRootForTest.view.captureRoboImage(file, roborazziOptions)
+      } else {
+        // For dialog
+//        windowRoots[1].decorView.rootView.captureRoboImage(file, roborazziOptions)
+//        captureScreenRoboImage()
+        captureRootsInternal(windowRoots.drop(1), roborazziOptions, file)
+      }
     }
 
     // Closing the activity is necessary to prevent memory leaks.
diff --git a/roborazzi/src/main/java/com/github/takahirom/roborazzi/Roborazzi.kt b/roborazzi/src/main/java/com/github/takahirom/roborazzi/Roborazzi.kt
index 336b0c126..ffe4d9207 100644
--- a/roborazzi/src/main/java/com/github/takahirom/roborazzi/Roborazzi.kt
+++ b/roborazzi/src/main/java/com/github/takahirom/roborazzi/Roborazzi.kt
@@ -144,18 +144,22 @@ fun captureScreenRoboImage(
   // Views needs to be laid out before we can capture them
   Espresso.onIdle()
 
-  val rootsOracle = RootsOracle_Factory({ Looper.getMainLooper() })
-    .get()
-  // Invoke rootOracle.listActiveRoots() via reflection
-  val listActiveRoots = rootsOracle.javaClass.getMethod("listActiveRoots")
-  listActiveRoots.isAccessible = true
-  @Suppress("UNCHECKED_CAST") val roots: List<Root> = listActiveRoots.invoke(rootsOracle) as List<Root>
+  val roots: List<Root> = fetchRobolectricWindowRoots()
   debugLog {
     "captureScreenRoboImage roots: ${roots.joinToString("\n") { it.toString() }}"
   }
+  captureRootsInternal(roots, roborazziOptions, file)
+}
+
+@InternalRoborazziApi
+fun captureRootsInternal(
+  roots: List<Root>,
+  roborazziOptions: RoborazziOptions,
+  file: File
+) {
   capture(
     rootComponent = RoboComponent.Screen(
-      rootsOrderByDepth = roots.sortedBy { it.windowLayoutParams.get()?.type },
+      rootsOrderByDepth = roots,
       roborazziOptions = roborazziOptions
     ),
     roborazziOptions = roborazziOptions,
@@ -169,6 +173,24 @@ fun captureScreenRoboImage(
   }
 }
 
+@InternalRoborazziApi
+fun fetchRobolectricWindowRoots(): List<Root> {
+  try {
+    @Suppress("INACCESSIBLE_TYPE") val rootsOracle = RootsOracle_Factory({ Looper.getMainLooper() })
+      .get()
+    // Invoke rootOracle.listActiveRoots() via reflection
+    val listActiveRoots = rootsOracle.javaClass.getMethod("listActiveRoots")
+    listActiveRoots.isAccessible = true
+    @Suppress("UNCHECKED_CAST") val roots: List<Root> =
+      (listActiveRoots.invoke(rootsOracle) as List<Root>
+        ).sortedBy { it.windowLayoutParams.get()?.type }
+    return roots
+  } catch(e: Throwable) {
+    e.printStackTrace()
+    return emptyList()
+  }
+}
+
 fun Bitmap.captureRoboImage(
   filePath: String = DefaultFileNameGenerator.generateFilePath(),
   roborazziOptions: RoborazziOptions = provideRoborazziContext().options,
diff --git a/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/ExampleUnitTest.kt b/sample-generate-preview-tests/src/test/java/com/github/takahirom/preview/ExampleUnitTest.kt
similarity index 79%
rename from sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/ExampleUnitTest.kt
rename to sample-generate-preview-tests/src/test/java/com/github/takahirom/preview/ExampleUnitTest.kt
index 6289244c0..9a10e100c 100644
--- a/sample-generate-preview-tests/src/test/java/com/github/takahirom/sample/ExampleUnitTest.kt
+++ b/sample-generate-preview-tests/src/test/java/com/github/takahirom/preview/ExampleUnitTest.kt
@@ -1,9 +1,8 @@
-package com.github.takahirom.sample
+package com.github.takahirom.preview
 
+import org.junit.Assert.assertEquals
 import org.junit.Test
 
-import org.junit.Assert.*
-
 /**
  * Example local unit test, which will execute on the development machine (host).
  *