diff --git a/ui/ui-test/src/androidTest/java/androidx/ui/test/SynchronizationMethodsTest.kt b/ui/ui-test/src/androidTest/java/androidx/ui/test/SynchronizationMethodsTest.kt
index a2c7d34..aa24a07f 100644
--- a/ui/ui-test/src/androidTest/java/androidx/ui/test/SynchronizationMethodsTest.kt
+++ b/ui/ui-test/src/androidTest/java/androidx/ui/test/SynchronizationMethodsTest.kt
@@ -17,6 +17,7 @@
 package androidx.ui.test
 
 import androidx.test.filters.MediumTest
+import androidx.ui.test.android.AndroidOwnerRegistry
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -35,7 +36,6 @@
     @Test
     fun runOnUiThread_void() {
         var called = false
-
         runOnUiThread { called = true }
         assertThat(called).isTrue()
     }
@@ -48,21 +48,35 @@
 
     @Test
     fun runOnIdleCompose() {
-        val result = runOnIdleCompose { "Hello" }
-        assertThat(result).isEqualTo("Hello")
+        withAndroidOwnerRegistry {
+            val result = runOnIdleCompose { "Hello" }
+            assertThat(result).isEqualTo("Hello")
+        }
     }
 
     @Test
     fun runOnIdleCompose_void() {
-        var called = false
-
-        runOnIdleCompose { called = true }
-        assertThat(called).isTrue()
+        withAndroidOwnerRegistry {
+            var called = false
+            runOnIdleCompose { called = true }
+            assertThat(called).isTrue()
+        }
     }
 
     @Test
     fun runOnIdleCompose_nullable() {
-        val result: String? = runOnIdleCompose { null }
-        assertThat(result).isEqualTo(null)
+        withAndroidOwnerRegistry {
+            val result: String? = runOnIdleCompose { null }
+            assertThat(result).isEqualTo(null)
+        }
+    }
+
+    private fun withAndroidOwnerRegistry(block: () -> Unit) {
+        AndroidOwnerRegistry.setupRegistry()
+        try {
+            block()
+        } finally {
+            AndroidOwnerRegistry.tearDownRegistry()
+        }
     }
 }
\ No newline at end of file
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/Assertions.kt b/ui/ui-test/src/main/java/androidx/ui/test/Assertions.kt
index 077b462..e9dbe65 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/Assertions.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/Assertions.kt
@@ -16,6 +16,7 @@
 
 package androidx.ui.test
 
+import androidx.ui.core.AndroidOwner
 import androidx.ui.core.LayoutNode
 import androidx.ui.core.findClosestParentNode
 import androidx.ui.core.semantics.SemanticsNode
@@ -23,9 +24,9 @@
 import androidx.ui.geometry.Rect
 import androidx.ui.semantics.AccessibilityRangeInfo
 import androidx.ui.semantics.SemanticsProperties
-import androidx.ui.test.android.SynchronizedTreeCollector
 import androidx.ui.unit.PxBounds
 import androidx.ui.unit.PxPosition
+import androidx.ui.unit.PxSize
 import androidx.ui.unit.height
 import androidx.ui.unit.px
 import androidx.ui.unit.toPx
@@ -296,30 +297,29 @@
 
     // check node doesn't clip unintentionally (e.g. row too small for content)
     val globalRect = node.globalBounds
-    if (!isInScreenBounds(globalRect)) {
+    if (!node.isInScreenBounds()) {
         return false
     }
 
     return (globalRect.width > 0.px && globalRect.height > 0.px)
 }
 
-internal fun isInScreenBounds(rectangle: PxBounds): Boolean {
-    if (rectangle.width == 0.px && rectangle.height == 0.px) {
+private fun SemanticsNode.isInScreenBounds(): Boolean {
+    val nodeBounds = globalBounds
+    if (nodeBounds.width == 0.px && nodeBounds.height == 0.px) {
         return false
     }
-    val displayMetrics = SynchronizedTreeCollector.collectOwners()
-        .findActivity()
-        .resources
-        .displayMetrics
 
-    val bottomRight = PxPosition(
-        displayMetrics.widthPixels.px,
-        displayMetrics.heightPixels.px
+    val displayMetrics = (componentNode.owner as AndroidOwner).view.resources.displayMetrics
+    val screenBounds = PxBounds(
+        PxPosition.Origin,
+        PxSize(displayMetrics.widthPixels.px, displayMetrics.heightPixels.px)
     )
-    return rectangle.top >= 0.px &&
-            rectangle.left >= 0.px &&
-            rectangle.right <= bottomRight.x &&
-            rectangle.bottom <= bottomRight.y
+
+    return nodeBounds.top >= screenBounds.top &&
+            nodeBounds.left >= screenBounds.left &&
+            nodeBounds.right <= screenBounds.right &&
+            nodeBounds.bottom <= screenBounds.bottom
 }
 
 /**
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/BitmapHelpers.kt b/ui/ui-test/src/main/java/androidx/ui/test/BitmapHelpers.kt
index b65642e..0173dce 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/BitmapHelpers.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/BitmapHelpers.kt
@@ -21,11 +21,8 @@
 import android.content.ContextWrapper
 import android.graphics.Bitmap
 import android.os.Build
-import android.os.Handler
-import android.os.Looper
 import android.view.View
 import androidx.annotation.RequiresApi
-import androidx.ui.core.semantics.SemanticsNode
 import androidx.ui.geometry.Offset
 import androidx.ui.geometry.Rect
 import androidx.ui.graphics.Canvas
@@ -35,7 +32,6 @@
 import androidx.ui.graphics.Shape
 import androidx.ui.graphics.addOutline
 import androidx.ui.graphics.asAndroidPath
-import androidx.ui.test.android.SynchronizedTreeCollector
 import androidx.ui.test.android.captureRegionToBitmap
 import androidx.ui.unit.Density
 import androidx.ui.unit.Dp
@@ -50,29 +46,6 @@
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertTrue
 
-@RequiresApi(Build.VERSION_CODES.O)
-internal fun captureNodeToBitmap(node: SemanticsNode): Bitmap {
-    val collectedInfo = SynchronizedTreeCollector.collectOwners()
-
-    val exists = getAllSemanticsNodes().any { it.id == node.id }
-    if (!exists) {
-        throw AssertionError("The required node is no longer in the tree!")
-    }
-
-    val window = collectedInfo.findActivity().window
-
-    // TODO(pavlis): Consider doing assertIsDisplayed here. Will need to move things around.
-
-    // TODO(pavlis): Make sure that the Activity actually hosts the view. As in case of popup
-    // it wouldn't. This will require us rewriting the structure how we collect the nodes.
-
-    // TODO(pavlis): Add support for popups. So if we find composable hosted in popup we can
-    // grab its reference to its window (need to add a hook to popup).
-
-    val handler = Handler(Looper.getMainLooper())
-    return captureRegionToBitmap(node.globalBounds.toRect(), handler, window)
-}
-
 /**
  * Captures the underlying component's surface into bitmap.
  *
@@ -82,8 +55,9 @@
  */
 @RequiresApi(Build.VERSION_CODES.O)
 fun SemanticsNodeInteraction.captureToBitmap(): Bitmap {
-    val errorMessageOnFail = "Failed to capture a node to bitmap."
-    return captureNodeToBitmap(fetchSemanticsNode(errorMessageOnFail))
+    val node = fetchSemanticsNode("Failed to capture a node to bitmap.")
+    // TODO(pavlis): Consider doing assertIsDisplayed here. Will need to move things around.
+    return captureRegionToBitmap(node.globalBounds.toRect(), node.componentNode.owner!!)
 }
 
 /**
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/Finders.kt b/ui/ui-test/src/main/java/androidx/ui/test/Finders.kt
index 1dcf1c8..66e4164 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/Finders.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/Finders.kt
@@ -91,5 +91,5 @@
 }
 
 internal fun getAllSemanticsNodes(): List<SemanticsNode> {
-    return SynchronizedTreeCollector.collectOwners().getAllSemanticNodes()
+    return SynchronizedTreeCollector.collectAllSemanticsNodes()
 }
\ No newline at end of file
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidInputDispatcher.kt b/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidInputDispatcher.kt
index bb03972..bf03018 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidInputDispatcher.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidInputDispatcher.kt
@@ -239,7 +239,7 @@
     }
 
     /**
-     * Sends the [event] to the [CollectedOwners] and [recycles][MotionEvent.recycle] it
+     * Sends the [event] to the MotionEvent dispatcher and [recycles][MotionEvent.recycle] it
      * regardless of the result. This method blocks until the event is sent.
      */
     private fun sendAndRecycleEvent(event: MotionEvent) {
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidOwnerRegistry.kt b/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidOwnerRegistry.kt
index d62da1b..44f9aa2 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidOwnerRegistry.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/android/AndroidOwnerRegistry.kt
@@ -37,7 +37,7 @@
     /**
      * Returns if the registry is setup to receive registrations from [AndroidOwner]s
      */
-    val isSetup: Boolean
+    val isSetUp: Boolean
         get() = AndroidOwner.onAndroidOwnerCreatedCallback == ::onAndroidOwnerCreated
 
     /**
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/android/SynchronizedTreeCollector.kt b/ui/ui-test/src/main/java/androidx/ui/test/android/SynchronizedTreeCollector.kt
index eecbc3b..5e95c72 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/android/SynchronizedTreeCollector.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/android/SynchronizedTreeCollector.kt
@@ -16,9 +16,6 @@
 
 package androidx.ui.test.android
 
-import android.app.Activity
-import android.content.Context
-import android.content.ContextWrapper
 import android.view.Choreographer
 import androidx.compose.onCommit
 import androidx.test.espresso.Espresso
@@ -31,31 +28,32 @@
 import java.util.concurrent.TimeUnit
 
 /**
- * Collects all [AndroidOwner]s that are part of the currently visible window.
+ * Collects all [SemanticsNode]s that are part of all compose hierarchies hosted by resumed
+ * activities.
  *
  * This operation is performed only after compose is idle via Espresso.
  */
 internal object SynchronizedTreeCollector {
     /**
-     * Collects all [AndroidOwner]s that are part of the currently visible window. Can only be
-     * used when using [ComposeTestRule][androidx.ui.test.ComposeTestRule]
+     * Collects all [SemanticsNode]s from all compose hierarchies. Can only be used when using
+     * [ComposeTestRule][androidx.ui.test.ComposeTestRule].
      *
      * This is a blocking call. Returns only after compose is idle.
      *
      * Can crash in case Espresso hits time out. This is not supposed to be handled as it
      * surfaces only in incorrect tests.
      */
-    internal fun collectOwners(): CollectedOwners {
-        check(AndroidOwnerRegistry.isSetup) {
-            "Test not setup properly. Use a ComposeTestRule in your test to be able to interact " +
-                    "with composables"
-        }
+    internal fun collectAllSemanticsNodes(): List<SemanticsNode> {
+        ensureAndroidOwnerRegistryIsSetUp()
+
+        // TODO(pavlis): Instead of returning a flatMap, let all consumers handle a tree
+        //  structure. In case of multiple AndroidOwners, add a fake root
         waitForIdle()
 
-        return CollectedOwners(AndroidOwnerRegistry.getOwners().also {
+        return AndroidOwnerRegistry.getOwners().also {
             // TODO(b/153632210): This check should be done by callers of collectOwners
             check(it.isNotEmpty()) { "No AndroidOwners found. Is your Activity resumed?" }
-        })
+        }.flatMap { it.semanticsOwner.getAllSemanticsNodes() }
     }
 
     /**
@@ -85,6 +83,13 @@
         waitForOnCommitCallbacks()
     }
 
+    private fun ensureAndroidOwnerRegistryIsSetUp() {
+        check(AndroidOwnerRegistry.isSetUp) {
+            "Test not setup properly. Use a ComposeTestRule in your test to be able to interact " +
+                    "with composables"
+        }
+    }
+
     /**
      * Waits for all scheduled [onCommit] callbacks to be executed.
      */
@@ -98,6 +103,8 @@
     }
 
     private fun waitForAndroidOwners() {
+        ensureAndroidOwnerRegistryIsSetUp()
+
         fun hasAndroidOwners(): Boolean = AndroidOwnerRegistry.getOwners().isNotEmpty()
 
         if (!hasAndroidOwners()) {
@@ -120,36 +127,3 @@
         }
     }
 }
-
-/**
- * There can be multiple Compose views in the Android hierarchy and we want to interact with all
- * of them. This class merges all the [AndroidOwner]s into one, hiding the fact that the API
- * might be interacting with several Compose roots.
- */
-internal data class CollectedOwners(val owners: Set<AndroidOwner>) {
-    // Recursively search for the Activity context through (possible) ContextWrappers
-    private fun Context.getActivity(): Activity? {
-        return when (this) {
-            is Activity -> this
-            is ContextWrapper -> this.baseContext.getActivity()
-            else -> null
-        }
-    }
-
-    fun findActivity(): Activity {
-        owners.forEach {
-            val activity = it.view.context.getActivity()
-            if (activity != null) {
-                return activity
-            }
-        }
-        throw AssertionError(
-            "Out of ${owners.size} Owners, none were attached to an Activity"
-        )
-    }
-
-    fun getAllSemanticNodes(): List<SemanticsNode> {
-        // TODO(pavlis): Once we have a tree support we will just add a fake root parent here
-        return owners.flatMap { it.semanticsOwner.getAllSemanticsNodes() }
-    }
-}
diff --git a/ui/ui-test/src/main/java/androidx/ui/test/android/WindowCapture.kt b/ui/ui-test/src/main/java/androidx/ui/test/android/WindowCapture.kt
index 7f5e5b9..caf7ee8 100644
--- a/ui/ui-test/src/main/java/androidx/ui/test/android/WindowCapture.kt
+++ b/ui/ui-test/src/main/java/androidx/ui/test/android/WindowCapture.kt
@@ -16,12 +16,18 @@
 
 package androidx.ui.test.android
 
+import android.app.Activity
+import android.content.Context
+import android.content.ContextWrapper
 import android.graphics.Bitmap
 import android.os.Build
 import android.os.Handler
+import android.os.Looper
 import android.view.PixelCopy
 import android.view.Window
 import androidx.annotation.RequiresApi
+import androidx.ui.core.AndroidOwner
+import androidx.ui.core.Owner
 import androidx.ui.geometry.Rect
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
@@ -29,21 +35,48 @@
 
 @RequiresApi(Build.VERSION_CODES.O)
 internal fun captureRegionToBitmap(
-    globalRect: Rect,
+    captureRect: Rect,
+    owner: Owner
+): Bitmap {
+
+    fun Context.getActivity(): Activity? {
+        return when (this) {
+            is Activity -> this
+            is ContextWrapper -> this.baseContext.getActivity()
+            else -> null
+        }
+    }
+
+    // TODO(pavlis): Make sure that the Activity actually hosts the view. As in case of popup
+    //  it wouldn't. This will require us rewriting the structure how we collect the nodes.
+
+    // TODO(pavlis): Add support for popups. So if we find composable hosted in popup we can
+    //  grab its reference to its window (need to add a hook to popup).
+
+    val window = (owner as AndroidOwner).view.context.getActivity()!!.window
+    val handler = Handler(Looper.getMainLooper())
+    return captureRegionToBitmap(captureRect, handler, window)
+}
+
+@RequiresApi(Build.VERSION_CODES.O)
+internal fun captureRegionToBitmap(
+    captureRect: Rect,
     handler: Handler,
     window: Window
 ): Bitmap {
     val destBitmap = Bitmap.createBitmap(
-        globalRect.width.roundToInt(),
-        globalRect.height.roundToInt(),
-        Bitmap.Config.ARGB_8888)
+        captureRect.width.roundToInt(),
+        captureRect.height.roundToInt(),
+        Bitmap.Config.ARGB_8888
+    )
 
     // TODO: This could go to some Android specific extensions.
     val srcRect = android.graphics.Rect(
-        globalRect.left.roundToInt(),
-        globalRect.top.roundToInt(),
-        globalRect.right.roundToInt(),
-        globalRect.bottom.roundToInt())
+        captureRect.left.roundToInt(),
+        captureRect.top.roundToInt(),
+        captureRect.right.roundToInt(),
+        captureRect.bottom.roundToInt()
+    )
 
     val latch = CountDownLatch(1)
     var copyResult = 0