diff --git a/MacrobenchmarkSample/app/build.gradle.kts b/MacrobenchmarkSample/app/build.gradle.kts
index 23beb2ca..8affb93a 100644
--- a/MacrobenchmarkSample/app/build.gradle.kts
+++ b/MacrobenchmarkSample/app/build.gradle.kts
@@ -22,7 +22,7 @@ plugins {
 }
 
 android {
-    compileSdk = 35
+    compileSdk = 36
     namespace = "com.example.macrobenchmark.target"
 
     defaultConfig {
@@ -30,7 +30,7 @@ android {
         versionCode = 1
         versionName = "1.0"
         minSdk = 24
-        targetSdk = 35
+        targetSdk = 36
         testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner"
     }
 
diff --git a/MacrobenchmarkSample/app/src/main/AndroidManifest.xml b/MacrobenchmarkSample/app/src/main/AndroidManifest.xml
index ab34acde..6536dfe1 100644
--- a/MacrobenchmarkSample/app/src/main/AndroidManifest.xml
+++ b/MacrobenchmarkSample/app/src/main/AndroidManifest.xml
@@ -14,7 +14,7 @@
         android:theme="@style/Theme.Macrobenchmarktarget">
         <profileable
             android:shell="true"
-            tools:targetApi="29,q" />
+            tools:targetApi="29" />
         <!-- [END macrobenchmark_setup_profileable] -->
         <!--
         Activities need to be exported so the Macrobenchmarks can discover them
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/FullyDrawnStartupActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/FullyDrawnStartupActivity.kt
index f6f37744..fdb1c64e 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/FullyDrawnStartupActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/FullyDrawnStartupActivity.kt
@@ -20,10 +20,13 @@ import android.os.Bundle
 import androidx.activity.ComponentActivity
 import androidx.activity.compose.ReportDrawnWhen
 import androidx.activity.compose.setContent
+import androidx.activity.enableEdgeToEdge
 import androidx.activity.viewModels
+import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
-import androidx.compose.material.CircularProgressIndicator
-import androidx.compose.material.Text
+import androidx.compose.material3.CircularProgressIndicator
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
@@ -37,20 +40,25 @@ class FullyDrawnStartupActivity : ComponentActivity() {
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
+        enableEdgeToEdge()
         setContent {
-            var isLoaded by remember { mutableStateOf(false) }
-            ReportDrawnWhen { isLoaded }
+            MaterialTheme {
+                Box {
+                    var isLoaded by remember { mutableStateOf(false) }
+                    ReportDrawnWhen { isLoaded }
 
-            LaunchedEffect(Unit) {
-                isLoaded = sampleViewModel.data.isReady()
-            }
+                    LaunchedEffect(Unit) {
+                        isLoaded = sampleViewModel.data.isReady()
+                    }
 
-            Column {
-                Text("Compose Macrobenchmark Target")
-                if (!isLoaded) {
-                    CircularProgressIndicator()
-                } else {
-                    Text("Fully Drawn")
+                    Column {
+                        Text("Compose Macrobenchmark Target")
+                        if (!isLoaded) {
+                            CircularProgressIndicator()
+                        } else {
+                            Text("Fully Drawn")
+                        }
+                    }
                 }
             }
         }
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/MainActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/MainActivity.kt
index 5c3d4236..c53bd52d 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/MainActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/MainActivity.kt
@@ -26,12 +26,12 @@ import androidx.compose.foundation.BorderStroke
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
-import androidx.compose.material.MaterialTheme
-import androidx.compose.material.Text
-import androidx.compose.material.TextButton
-import androidx.compose.material.TopAppBar
+import androidx.compose.material3.ExperimentalMaterial3Api
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
+import androidx.compose.material3.TextButton
+import androidx.compose.material3.TopAppBar
 import androidx.compose.runtime.Composable
-import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.semantics.semantics
@@ -57,7 +57,7 @@ private const val TAG = "MainActivity"
 
 class MainActivity : ComponentActivity() {
 
-    @OptIn(ExperimentalComposeUiApi::class)
+    @OptIn(ExperimentalMaterial3Api::class)
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         setContent {
@@ -132,7 +132,7 @@ class MainActivity : ComponentActivity() {
                 .padding(8.dp)
                 .testTag(testTag),
             onClick = onClick,
-            border = BorderStroke(1.dp, MaterialTheme.colors.primary)
+            border = BorderStroke(1.dp, MaterialTheme.colorScheme.primary)
         ) {
             Text(name)
         }
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ComposeActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ComposeActivity.kt
index cbef2597..a1943ff9 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ComposeActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ComposeActivity.kt
@@ -21,7 +21,6 @@ import androidx.activity.ComponentActivity
 import androidx.activity.compose.setContent
 import androidx.appcompat.app.AlertDialog
 import androidx.compose.foundation.clickable
-import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
@@ -30,11 +29,13 @@ import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.wrapContentSize
 import androidx.compose.foundation.lazy.LazyColumn
 import androidx.compose.foundation.lazy.items
-import androidx.compose.material.Card
-import androidx.compose.material.Checkbox
-import androidx.compose.material.MaterialTheme
-import androidx.compose.material.Text
-import androidx.compose.material.TextField
+import androidx.compose.material3.Card
+import androidx.compose.material3.Checkbox
+import androidx.compose.material3.ExperimentalMaterial3Api
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
+import androidx.compose.material3.TextField
+import androidx.compose.material3.TopAppBar
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
@@ -43,16 +44,12 @@ import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.platform.testTag
-import androidx.compose.ui.semantics.semantics
-import androidx.compose.ui.semantics.testTagsAsResourceId
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.dp
 import androidx.tracing.trace
 import com.example.macrobenchmark.target.recyclerview.Entry
 import com.example.macrobenchmark.target.util.ClickTrace
 
-@OptIn(ExperimentalComposeUiApi::class)
 class ComposeActivity : ComponentActivity() {
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
@@ -68,49 +65,35 @@ class ComposeActivity : ComponentActivity() {
         }
     }
 
+    @OptIn(ExperimentalMaterial3Api::class)
     @Composable
     private fun BenchmarkComposeList(data: List<Entry>) {
         MaterialTheme {
-            Box(
-                modifier = Modifier.semantics {
-                    // Allows to use testTag() for UiAutomator's resource-id.
-                    // It can be enabled high in the compose hierarchy,
-                    // so that it's enabled for the whole subtree
-                    testTagsAsResourceId = true
-                }
-            ) {
-                // Thanks to [SemanticsPropertyReceiver.testTagsAsResourceId],
-                // [Modifier.testTag]s will be propagated to resource-id
-                // and can be accessed from benchmarks.
-                var value by remember { mutableStateOf("Enter text here") }
-                Column(modifier = Modifier.fillMaxWidth()) {
-                    TextField(
-                        modifier = Modifier
-                            .fillMaxWidth()
-                            .testTag("input"),
-                        value = value,
-                        onValueChange = { value = it },
-                        placeholder = { Text("Enter text here") }
-                    )
+            var value by remember { mutableStateOf("Enter text here") }
+            Column(modifier = Modifier.fillMaxWidth()) {
+                TopAppBar(title = { Text("Compose Sample") })
+                TextField(
+                    modifier = Modifier
+                        .fillMaxWidth(),
+                    value = value,
+                    onValueChange = { value = it },
+                    placeholder = { Text("Enter text here") }
+                )
 
-                    LazyColumn(
-                        modifier = Modifier
-                            .testTag("myLazyColumn")
-                    ) {
-                        items(data, key = { it.contents }) { item ->
-                            EntryRow(
-                                entry = item,
-                                modifier = Modifier
-                                    .padding(8.dp)
-                                    .clickable {
-                                        ClickTrace.onClickPerformed()
-                                        AlertDialog
-                                            .Builder(this@ComposeActivity)
-                                            .setMessage("Item clicked")
-                                            .show()
-                                    }
-                            )
-                        }
+                LazyColumn {
+                    items(data, key = { it.contents }) { item ->
+                        EntryRow(
+                            entry = item,
+                            modifier = Modifier
+                                .padding(8.dp)
+                                .clickable {
+                                    ClickTrace.onClickPerformed()
+                                    AlertDialog
+                                        .Builder(this@ComposeActivity)
+                                        .setMessage("Item clicked")
+                                        .show()
+                                }
+                        )
                     }
                 }
             }
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ListViewActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ListViewActivity.kt
index b1586515..a5b16588 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ListViewActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ListViewActivity.kt
@@ -24,6 +24,9 @@ import android.widget.BaseAdapter
 import android.widget.TextView
 import androidx.appcompat.app.AlertDialog
 import androidx.appcompat.app.AppCompatActivity
+import androidx.core.view.ViewCompat
+import androidx.core.view.WindowInsetsCompat
+import androidx.core.view.updatePadding
 import com.example.macrobenchmark.target.databinding.ActivityListViewBinding
 import com.example.macrobenchmark.target.recyclerview.Entry
 import com.example.macrobenchmark.target.util.ClickTrace
@@ -73,5 +76,19 @@ class ListViewActivity : AppCompatActivity() {
                 .show()
         }
 
+        ViewCompat.setOnApplyWindowInsetsListener(binding.listview) { v, insets ->
+            val bars = insets.getInsets(
+                WindowInsetsCompat.Type.systemBars()
+                        or WindowInsetsCompat.Type.displayCutout()
+            )
+            v.updatePadding(
+                left = bars.left,
+                top = bars.top,
+                right = bars.right,
+                bottom = bars.bottom,
+            )
+            WindowInsetsCompat.CONSUMED
+        }
+
     }
 }
\ No newline at end of file
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/NestedRecyclerActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/NestedRecyclerActivity.kt
index 12905ded..299b0fd5 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/NestedRecyclerActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/NestedRecyclerActivity.kt
@@ -19,6 +19,9 @@ package com.example.macrobenchmark.target.activity.clicklatency
 import android.os.Bundle
 import androidx.activity.viewModels
 import androidx.appcompat.app.AppCompatActivity
+import androidx.core.view.ViewCompat
+import androidx.core.view.WindowInsetsCompat
+import androidx.core.view.updatePadding
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
@@ -55,6 +58,20 @@ class NestedRecyclerActivity : AppCompatActivity() {
                 }
             }
         }
+
+        ViewCompat.setOnApplyWindowInsetsListener(binding.recycler) { v, insets ->
+            val bars = insets.getInsets(
+                WindowInsetsCompat.Type.systemBars()
+                        or WindowInsetsCompat.Type.displayCutout()
+            )
+            v.updatePadding(
+                left = bars.left,
+                top = bars.top,
+                right = bars.right,
+                bottom = bars.bottom,
+            )
+            WindowInsetsCompat.CONSUMED
+        }
     }
 }
 
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/RecyclerViewActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/RecyclerViewActivity.kt
index dd33aae6..d2e61f69 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/RecyclerViewActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/RecyclerViewActivity.kt
@@ -18,6 +18,9 @@ package com.example.macrobenchmark.target.activity.clicklatency
 
 import android.os.Bundle
 import androidx.appcompat.app.AppCompatActivity
+import androidx.core.view.ViewCompat
+import androidx.core.view.WindowInsetsCompat
+import androidx.core.view.updatePadding
 import androidx.recyclerview.widget.LinearLayoutManager
 import com.example.macrobenchmark.target.databinding.ActivityRecyclerViewBinding
 import com.example.macrobenchmark.target.recyclerview.Entry
@@ -30,6 +33,20 @@ open class RecyclerViewActivity : AppCompatActivity() {
         val binding = ActivityRecyclerViewBinding.inflate(layoutInflater)
         setContentView(binding.root)
 
+        ViewCompat.setOnApplyWindowInsetsListener(binding.recycler) { v, insets ->
+            val bars = insets.getInsets(
+                WindowInsetsCompat.Type.systemBars()
+                        or WindowInsetsCompat.Type.displayCutout()
+            )
+            v.updatePadding(
+                left = bars.left,
+                top = bars.top,
+                right = bars.right,
+                bottom = bars.bottom,
+            )
+            WindowInsetsCompat.CONSUMED
+        }
+
         // This argument allows the Macrobenchmark tests control the content being tested.
         // In your app, you could use this approach to navigate to a consistent UI.
         // e.g. Here the UI is being populated with a well known number of list items.
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ScrollViewActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ScrollViewActivity.kt
index 2e78d89d..94552baa 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ScrollViewActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/clicklatency/ScrollViewActivity.kt
@@ -21,6 +21,9 @@ import android.view.LayoutInflater
 import android.widget.TextView
 import androidx.appcompat.app.AlertDialog
 import androidx.appcompat.app.AppCompatActivity
+import androidx.core.view.ViewCompat
+import androidx.core.view.WindowInsetsCompat
+import androidx.core.view.updatePadding
 import com.example.macrobenchmark.target.R
 import com.example.macrobenchmark.target.databinding.ActivityScrollViewBinding
 import com.example.macrobenchmark.target.recyclerview.Entry
@@ -36,6 +39,21 @@ class ScrollViewActivity : AppCompatActivity() {
         title = "ScrollView Sample"
         val binding = ActivityScrollViewBinding.inflate(layoutInflater)
         setContentView(binding.root)
+
+        ViewCompat.setOnApplyWindowInsetsListener(binding.scrollcontent) { v, insets ->
+            val bars = insets.getInsets(
+                WindowInsetsCompat.Type.systemBars()
+                        or WindowInsetsCompat.Type.displayCutout()
+            )
+            v.updatePadding(
+                left = bars.left,
+                top = bars.top,
+                right = bars.right,
+                bottom = bars.bottom,
+            )
+            WindowInsetsCompat.CONSUMED
+        }
+
         val itemCount = intent.getIntExtra(RecyclerViewActivity.EXTRA_ITEM_COUNT, 1000)
 
         val items = List(itemCount) {
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/login/LoginActivity.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/login/LoginActivity.kt
index dba1029b..0205fc31 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/login/LoginActivity.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/activity/login/LoginActivity.kt
@@ -26,13 +26,15 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.safeDrawingPadding
 import androidx.compose.foundation.text.KeyboardOptions
-import androidx.compose.material.Button
-import androidx.compose.material.MaterialTheme
-import androidx.compose.material.Surface
-import androidx.compose.material.Text
-import androidx.compose.material.TextField
-import androidx.compose.material.TopAppBar
+import androidx.compose.material3.Button
+import androidx.compose.material3.ExperimentalMaterial3Api
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Surface
+import androidx.compose.material3.Text
+import androidx.compose.material3.TextField
+import androidx.compose.material3.TopAppBar
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
@@ -53,7 +55,6 @@ import com.example.macrobenchmark.target.activity.MainActivity
 import com.example.macrobenchmark.target.util.SampleViewModel
 import kotlinx.coroutines.launch
 
-@OptIn(ExperimentalComposeUiApi::class)
 class LoginActivity : ComponentActivity() {
 
     private val sampleViewModel by viewModels<SampleViewModel>()
@@ -68,10 +69,13 @@ class LoginActivity : ComponentActivity() {
             }
         }
         setContent {
-            LoginScreen()
+            MaterialTheme {
+                LoginScreen()
+            }
         }
     }
 
+    @OptIn(ExperimentalMaterial3Api::class)
     @Composable
     fun LoginScreen() {
         var userName by remember {
@@ -82,12 +86,14 @@ class LoginActivity : ComponentActivity() {
         }
         MaterialTheme {
             Box(
-                modifier = Modifier.semantics {
-                    // Allows to use testTag() for UiAutomator's resource-id.
-                    // It can be enabled high in the compose hierarchy,
-                    // so that it's enabled for the whole subtree
-                    testTagsAsResourceId = true
-                }
+                modifier = Modifier
+                    .semantics {
+                        // Allows to use testTag() for UiAutomator's resource-id.
+                        // It can be enabled high in the compose hierarchy,
+                        // so that it's enabled for the whole subtree
+                        testTagsAsResourceId = true
+                    }
+                    .safeDrawingPadding()
             ) {
 
                 Column(
diff --git a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/recyclerview/ParentAdapter.kt b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/recyclerview/ParentAdapter.kt
index fcaaa413..b691e1eb 100644
--- a/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/recyclerview/ParentAdapter.kt
+++ b/MacrobenchmarkSample/app/src/main/kotlin/com/example/macrobenchmark/target/recyclerview/ParentAdapter.kt
@@ -95,7 +95,7 @@ class ParentAdapter(
     override fun onViewRecycled(holder: BindingViewHolder<ItemParentBinding>) {
         trace("ParentAdapter.onViewRecycled") {
             // save the state of nested recycler, so we can show the same scrolled position
-            nestedStates[holder.adapterPosition] =
+            nestedStates[holder.absoluteAdapterPosition] =
                 holder.binding.rowRecycler.layoutManager?.onSaveInstanceState()
         }
         super.onViewRecycled(holder)
diff --git a/MacrobenchmarkSample/gradle/libs.versions.toml b/MacrobenchmarkSample/gradle/libs.versions.toml
index a7612762..8c438901 100644
--- a/MacrobenchmarkSample/gradle/libs.versions.toml
+++ b/MacrobenchmarkSample/gradle/libs.versions.toml
@@ -2,7 +2,7 @@
 agp = "8.11.1"
 activity = "1.10.1"
 appcompat = "1.7.1"
-benchmark = "1.3.4"
+benchmark = "1.4.0"
 composeBom = "2025.07.00"
 constraintLayout = "2.2.1"
 core = "1.16.0"
@@ -14,12 +14,13 @@ espressoCore = "3.6.1"
 jUnit = "1.2.1"
 kotlin = "2.2.0"
 lifecycle = "2.9.2"
-material = "1.12.0"
+material = "1.13.0-rc01"
+material3 = "1.4.0-beta01"
 profileInstaller = "1.4.1"
 runtimeTracing = "1.8.3"
 tracing = "1.3.0"
 tracingPerfetto = "1.0.0"
-uiAutomator = "2.3.0"
+uiAutomator = "2.4.0-alpha05"
 
 [libraries]
 compose-bom = { group = "androidx.compose", name = "compose-bom", version.ref = "composeBom" }
@@ -28,7 +29,7 @@ androidx-junit = { group = "androidx.test.ext", name = "junit", version.ref = "j
 appcompat = { group = "androidx.appcompat", name = "appcompat", version.ref = "appcompat" }
 benchmark-junit = { group = "androidx.benchmark", name = "benchmark-macro-junit4", version.ref = "benchmark" }
 compose-activity = { group = "androidx.activity", name = "activity-compose", version.ref = "activity" }
-compose-material = { group = "androidx.compose.material", name = "material" }
+compose-material = { group = "androidx.compose.material3", name = "material3" }
 compose-runtime-tracing = { module = "androidx.compose.runtime:runtime-tracing", version.ref = "runtimeTracing" }
 compose-ui = { group = "androidx.compose.ui", name = "ui" }
 compose-ui-tooling = { group = "androidx.compose.ui", name = "ui-tooling" }
diff --git a/MacrobenchmarkSample/macrobenchmark/build.gradle.kts b/MacrobenchmarkSample/macrobenchmark/build.gradle.kts
index a3c0764c..77371f60 100644
--- a/MacrobenchmarkSample/macrobenchmark/build.gradle.kts
+++ b/MacrobenchmarkSample/macrobenchmark/build.gradle.kts
@@ -26,14 +26,14 @@ plugins {
 // [START macrobenchmark_setup_android]
 android {
     // [START_EXCLUDE]
-    compileSdk = 35
+    compileSdk = 36
     namespace = "com.example.macrobenchmark"
 
     defaultConfig {
         // Minimum supported version for Baseline Profiles.
         // On lower APIs, apps are fully AOT compile, therefore Baseline Profiles aren't needed.
         minSdk = 28
-        targetSdk = 35
+        targetSdk = 36
         testInstrumentationRunner = "androidx.test.runner.AndroidJUnitRunner"
     }
 
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/Common.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/Common.kt
deleted file mode 100644
index 9738eafc..00000000
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/Common.kt
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * 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 com.example.macrobenchmark.baselineprofile
-
-import androidx.benchmark.macro.MacrobenchmarkScope
-import androidx.test.uiautomator.UiDevice
-
-const val TARGET_PACKAGE = "com.example.macrobenchmark.target"
-
-/**
- * Clears the application data for the package specified in the [MacrobenchmarkScope].
- * @param scope The [MacrobenchmarkScope] providing information about the benchmark,
- * including the package name of the app under test.
- */
-fun UiDevice.clearData(scope: MacrobenchmarkScope) {
-    val command = "pm clear ${scope.packageName}"
-    val output = executeShellCommand(command)
-//    Assert.assertEquals("Success", output)
-}
\ No newline at end of file
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/ComposeActivityBaselineProfileGenerator.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/ComposeActivityBaselineProfileGenerator.kt
index 7264400b..db1103d9 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/ComposeActivityBaselineProfileGenerator.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/ComposeActivityBaselineProfileGenerator.kt
@@ -19,10 +19,9 @@ package com.example.macrobenchmark.baselineprofile
 import android.content.Intent
 import androidx.benchmark.macro.junit4.BaselineProfileRule
 import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
-import com.example.macrobenchmark.benchmark.util.findOrFail
-import com.example.macrobenchmark.benchmark.util.waitAndFind
+import androidx.test.uiautomator.uiAutomator
+import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
@@ -42,16 +41,13 @@ class ComposeActivityBaselineProfileGenerator {
             maxIterations = 15,
             stableIterations = 3
         ) {
-            // Start into the Compose Activity
-            startActivityAndWait(Intent("$TARGET_PACKAGE.COMPOSE_ACTIVITY"))
-
-            // Scrolling through the Compose journey
-            device.waitAndFind(By.res("myLazyColumn")).also {
-                it.setGestureMargin(device.displayWidth / 10)
-                it.fling(Direction.DOWN)
+            uiAutomator {
+                // Start into the Compose Activity
+                startIntent(Intent("$TARGET_PACKAGE.COMPOSE_ACTIVITY"))
+                // Scrolling through the Compose journey
+                    onElement { viewIdResourceName == "myLazyColumn" }.fling(Direction.DOWN)
+                onElement { viewIdResourceName == "myLazyColumn" }.fling(Direction.UP)
             }
-
-            device.findOrFail(By.res("myLazyColumn")).fling(Direction.UP)
         }
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/LoginBaselineProfileGenerator.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/LoginBaselineProfileGenerator.kt
index 9792a33d..bcf2621e 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/LoginBaselineProfileGenerator.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/LoginBaselineProfileGenerator.kt
@@ -19,7 +19,9 @@ package com.example.macrobenchmark.baselineprofile
 import android.content.Intent
 import androidx.benchmark.macro.junit4.BaselineProfileRule
 import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
-import androidx.test.uiautomator.By
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
+import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -37,12 +39,13 @@ class LoginBaselineProfileGenerator {
             maxIterations = 15,
             stableIterations = 3
         ) {
-            device.clearData(this)
-            startActivityAndWait(Intent("$packageName.LOGIN_ACTIVITY"))
-            device.findObject(By.res("userName")).text = "user"
-            device.findObject(By.res("password")).text = "password"
-            device.findObject(By.res("login")).click()
-            device.waitForIdle()
+            uiAutomator {
+                clearAppData()
+                startIntent(Intent("$packageName.LOGIN_ACTIVITY"))
+                onElement { textAsString() == "User: " }.text = "user"
+                onElement { textAsString() == "Password: " }.text = "password"
+                onElement { textAsString() == "Login" }.click()
+            }
         }
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/RecyclerViewActivityBaselineProfileGenerator.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/RecyclerViewActivityBaselineProfileGenerator.kt
index a5ab1b8a..8ef08690 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/RecyclerViewActivityBaselineProfileGenerator.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/RecyclerViewActivityBaselineProfileGenerator.kt
@@ -19,8 +19,9 @@ package com.example.macrobenchmark.baselineprofile
 import android.content.Intent
 import androidx.benchmark.macro.junit4.BaselineProfileRule
 import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
+import androidx.test.uiautomator.uiAutomator
+import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
@@ -40,14 +41,13 @@ class RecyclerViewActivityBaselineProfileGenerator {
             maxIterations = 15,
             stableIterations = 3
         ) {
-            // Start into the RecyclerViewActivity
-            startActivityAndWait(Intent("$TARGET_PACKAGE.RECYCLER_VIEW_ACTIVITY"))
+            uiAutomator {
+                // Start into the RecyclerViewActivity
+                startIntent(Intent("$TARGET_PACKAGE.RECYCLER_VIEW_ACTIVITY"))
 
-            // Scrolling RecyclerView journey
-            device.findObject(By.res(packageName, "recycler")).also {
-                it.setGestureMargin(device.displayWidth / 10)
-                it.fling(Direction.DOWN)
-                it.fling(Direction.UP)
+                // Scrolling RecyclerView journey
+                onElement { viewIdResourceName == "recycler" }.fling(Direction.DOWN)
+                onElement { viewIdResourceName == "recycler" }.fling(Direction.UP)
             }
         }
     }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupProfileGenerator.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupProfileGenerator.kt
index ba884bbc..45a7ca97 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupProfileGenerator.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupProfileGenerator.kt
@@ -18,6 +18,8 @@ package com.example.macrobenchmark.baselineprofile
 
 import androidx.benchmark.macro.junit4.BaselineProfileRule
 import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
+import androidx.test.uiautomator.uiAutomator
+import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -41,7 +43,9 @@ class StartupProfileGenerator {
             stableIterations = 3,
             includeInStartupProfile = true
         ) {
-            startActivityAndWait()
+            uiAutomator {
+                startApp(TARGET_PACKAGE)
+            }
         }
     }
 }
\ No newline at end of file
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/LoginBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/LoginBenchmark.kt
index e94c0af7..8c22db90 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/LoginBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/LoginBenchmark.kt
@@ -25,10 +25,12 @@ import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.uiautomator.By
+import androidx.test.uiautomator.UiAutomatorTestScope
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
+import androidx.test.uiautomator.watcher.PermissionDialog
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
-import com.example.macrobenchmark.benchmark.permissions.allowNotifications
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -45,41 +47,59 @@ class LoginBenchmark {
             putString("user", "benchUser")
             putString("password", "benchPassword")
         }
-        benchmarkLoginActivity(extras = extras)
+        benchmarkLoginActivity(
+            extras = extras,
+            measureBlock = {
+                uiAutomator {
+                    clickOnLogin()
+                }
+            }
+        )
     }
 
     @Test
     fun loginInSetupBlock() {
-        benchmarkLoginActivity(setupBlock = {
-            startActivityAndWait(Intent("$packageName.LOGIN_ACTIVITY"))
-            login()
+        benchmarkLoginActivity(measureBlock = {
+            uiAutomator {
+                inputLoginDetails()
+                clickOnLogin()
+            }
         })
     }
 
     @Test
     fun loginWithUiAutomator() {
-        benchmarkLoginActivity {
-            login()
-        }
+        benchmarkLoginActivity(
+            measureBlock = {
+                uiAutomator {
+                    inputLoginDetails()
+                    clickOnLogin()
+                }
+            }
+        )
     }
 
     @Test
     fun loginInAfterPermissionsGranted() {
-        benchmarkLoginActivity(setupBlock = {
-            allowNotifications()
-
-            startActivityAndWait(Intent("$packageName.LOGIN_ACTIVITY"))
-            login()
+        benchmarkLoginActivity(measureBlock = {
+            uiAutomator {
+                watchFor(PermissionDialog) {
+                    clickAllow()
+                }
+                inputLoginDetails()
+                clickOnLogin()
+            }
         })
     }
 
-    private fun MacrobenchmarkScope.login() {
-        device.findObject(By.res("userName")).text = "user"
-        device.findObject(By.res("password")).text = "password"
-        device.findObject(By.res("login")).click()
-        device.waitForIdle()
+    private fun UiAutomatorTestScope.inputLoginDetails() {
+        onElement { viewIdResourceName == "userName" }.text = "user"
+        onElement { viewIdResourceName == "password" }.text = "password"
     }
 
+    private fun UiAutomatorTestScope.clickOnLogin() =
+        onElement { textAsString() == "Login" }.click()
+
     private fun benchmarkLoginActivity(
         extras: Bundle = Bundle(),
         setupBlock: MacrobenchmarkScope.() -> Unit = {},
@@ -93,11 +113,13 @@ class LoginBenchmark {
             iterations = DEFAULT_ITERATIONS,
             setupBlock = setupBlock,
         ) {
-            startActivityAndWait(
-                Intent()
-                    .putExtras(extras)
-                    .setAction("$packageName.LOGIN_ACTIVITY")
-            )
+            uiAutomator {
+                startIntent(
+                    Intent()
+                        .putExtras(extras)
+                        .setAction("$packageName.LOGIN_ACTIVITY")
+                )
+            }
             measureBlock()
         }
     }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/clickslatency/ClickLatencyBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/clickslatency/ClickLatencyBenchmark.kt
index 30b55ec6..3a81a575 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/clickslatency/ClickLatencyBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/clickslatency/ClickLatencyBenchmark.kt
@@ -19,16 +19,15 @@ package com.example.macrobenchmark.benchmark.clickslatency;
 import android.content.Intent
 import androidx.benchmark.macro.CompilationMode
 import androidx.benchmark.macro.ExperimentalMetricApi
-import androidx.benchmark.macro.MacrobenchmarkScope
 import androidx.benchmark.macro.TraceSectionMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
-import androidx.test.uiautomator.Until
+import androidx.test.uiautomator.UiAutomatorTestScope
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
-import junit.framework.TestCase.fail
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -49,22 +48,25 @@ class ClickLatencyBenchmark {
     fun simpleViewClick() {
         var firstStart = true
         benchmarkRule.measureRepeated(
-                packageName = TARGET_PACKAGE,
-                metrics = listOf(TraceSectionMetric("ClickTrace")),
-                compilationMode = CompilationMode.Full(),
-                startupMode = null,
-                iterations = DEFAULT_ITERATIONS,
-                setupBlock = {
-                    if (firstStart) {
-                        startActivityAndWait()
-                        firstStart = false
+            packageName = TARGET_PACKAGE,
+            metrics = listOf(TraceSectionMetric("ClickTrace")),
+            compilationMode = CompilationMode.Full(),
+            startupMode = null,
+            iterations = DEFAULT_ITERATIONS,
+            setupBlock = {
+                if (firstStart) {
+                    uiAutomator {
+                        startApp(TARGET_PACKAGE)
                     }
+                    firstStart = false
                 }
+            }
         ) {
-            clickOnId("launchRecyclerActivity")
-            waitForTextShown("RecyclerView Sample")
-            device.pressBack()
-            waitForTextGone("RecyclerView Sample")
+            uiAutomator {
+                onElement { viewIdResourceName == "launchRecyclerActivity" }.click()
+                onElement { text == "RecyclerView Sample" && isVisibleToUser }
+                pressBack()
+            }
         }
     }
 
@@ -72,20 +74,23 @@ class ClickLatencyBenchmark {
     fun recyclerViewClick() {
         var firstStart = true
         benchmarkRule.measureRepeated(
-                packageName = TARGET_PACKAGE,
-                metrics = listOf(TraceSectionMetric("ClickTrace")),
-                compilationMode = CompilationMode.Full(),
-                startupMode = null,
-                iterations = DEFAULT_ITERATIONS,
-                setupBlock = {
-                    if (firstStart) {
-                        val intent = Intent("$packageName.RECYCLER_VIEW_ACTIVITY")
-                        startActivityAndWait(intent)
-                        firstStart = false
+            packageName = TARGET_PACKAGE,
+            metrics = listOf(TraceSectionMetric("ClickTrace")),
+            compilationMode = CompilationMode.Full(),
+            startupMode = null,
+            iterations = DEFAULT_ITERATIONS,
+            setupBlock = {
+                if (firstStart) {
+                    uiAutomator {
+                        startIntent(Intent("$packageName.RECYCLER_VIEW_ACTIVITY"))
                     }
+                    firstStart = false
                 }
+            }
         ) {
-            clickOnFirstItem()
+            uiAutomator {
+                clickOnFirstItem()
+            }
         }
     }
 
@@ -93,20 +98,23 @@ class ClickLatencyBenchmark {
     fun composeLazyColumnClick() {
         var firstStart = true
         benchmarkRule.measureRepeated(
-                packageName = TARGET_PACKAGE,
-                metrics = listOf(TraceSectionMetric("ClickTrace")),
-                compilationMode = CompilationMode.Full(),
-                startupMode = null,
-                iterations = DEFAULT_ITERATIONS,
-                setupBlock = {
-                    if (firstStart) {
-                        val intent = Intent("$packageName.COMPOSE_ACTIVITY")
-                        startActivityAndWait(intent)
+            packageName = TARGET_PACKAGE,
+            metrics = listOf(TraceSectionMetric("ClickTrace")),
+            compilationMode = CompilationMode.Full(),
+            startupMode = null,
+            iterations = DEFAULT_ITERATIONS,
+            setupBlock = {
+                if (firstStart) {
+                    uiAutomator {
+                        startIntent(Intent("$packageName.COMPOSE_ACTIVITY"))
                         firstStart = false
                     }
                 }
+            }
         ) {
-            clickOnFirstItem()
+            uiAutomator {
+                clickOnFirstItem()
+            }
         }
     }
 
@@ -114,20 +122,23 @@ class ClickLatencyBenchmark {
     fun listViewClick() {
         var firstStart = true
         benchmarkRule.measureRepeated(
-                packageName = TARGET_PACKAGE,
-                metrics = listOf(TraceSectionMetric("ClickTrace")),
-                compilationMode = CompilationMode.Full(),
-                startupMode = null,
-                iterations = DEFAULT_ITERATIONS,
-                setupBlock = {
-                    if (firstStart) {
-                        val intent = Intent("$packageName.LIST_VIEW_ACTIVITY")
-                        startActivityAndWait(intent)
+            packageName = TARGET_PACKAGE,
+            metrics = listOf(TraceSectionMetric("ClickTrace")),
+            compilationMode = CompilationMode.Full(),
+            startupMode = null,
+            iterations = DEFAULT_ITERATIONS,
+            setupBlock = {
+                if (firstStart) {
+                    uiAutomator {
+                        startIntent(Intent("$packageName.LIST_VIEW_ACTIVITY"))
                         firstStart = false
                     }
                 }
+            }
         ) {
-            clickOnFirstItem()
+            uiAutomator {
+                clickOnFirstItem()
+            }
         }
     }
 
@@ -135,59 +146,30 @@ class ClickLatencyBenchmark {
     fun scrollViewClick() {
         var firstStart = true
         benchmarkRule.measureRepeated(
-                packageName = TARGET_PACKAGE,
-                metrics = listOf(TraceSectionMetric("ClickTrace")),
-                compilationMode = CompilationMode.Full(),
-                startupMode = null,
-                iterations = DEFAULT_ITERATIONS,
-                setupBlock = {
-                    if (firstStart) {
-                        val intent = Intent("$packageName.SCROLL_VIEW_ACTIVITY")
-                        startActivityAndWait(intent)
+            packageName = TARGET_PACKAGE,
+            metrics = listOf(TraceSectionMetric("ClickTrace")),
+            compilationMode = CompilationMode.Full(),
+            startupMode = null,
+            iterations = DEFAULT_ITERATIONS,
+            setupBlock = {
+                if (firstStart) {
+                    uiAutomator {
+                        startIntent(Intent("$packageName.SCROLL_VIEW_ACTIVITY"))
                         firstStart = false
                     }
                 }
+            }
         ) {
-            clickOnFirstItem()
+            uiAutomator {
+                clickOnFirstItem()
+            }
         }
     }
 
-    private fun MacrobenchmarkScope.clickOnFirstItem() {
-        clickOnText("Item 0")
-        waitForTextShown("Item clicked")
+    private fun UiAutomatorTestScope.clickOnFirstItem() {
+        onElement { textAsString() == "Item 0" }.click()
+        onElement { textAsString() == "Item clicked" && isVisibleToUser }
         // Dismiss dialog
-        device.pressBack()
-        waitForTextGone("Item clicked")
-    }
-
-    private fun MacrobenchmarkScope.waitForTextShown(text: String) {
-        check(device.wait(Until.hasObject(By.text(text)), 500)) {
-            "View showing '$text' not found after waiting 500 ms."
-        }
-    }
-
-    private fun MacrobenchmarkScope.waitForTextGone(text: String) {
-        check(device.wait(Until.gone(By.text(text)), 500)) {
-            "View showing '$text' not found after waiting 500 ms."
-        }
-    }
-
-    private fun MacrobenchmarkScope.clickOnText(text: String) {
-        device
-            .findObject(By.text(text))
-            .click()
-    }
-
-    private fun MacrobenchmarkScope.clickOnId(resourceId: String) {
-        val selector = By.res(resourceId)
-        if (!device.wait(Until.hasObject(selector), 2_500)) {
-            fail("Did not find object with id $resourceId")
-        }
-
-        device
-            .findObject(selector)
-            .click()
-        // Chill to ensure we capture the end of the click span in the trace.
-        Thread.sleep(100)
+        pressBack()
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/FrameTimingBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/FrameTimingBenchmark.kt
index 74abacf9..a21009b9 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/FrameTimingBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/FrameTimingBenchmark.kt
@@ -17,7 +17,6 @@
 package com.example.macrobenchmark.benchmark.frames
 
 import android.content.Intent
-import android.graphics.Point
 import androidx.benchmark.macro.CompilationMode
 import androidx.benchmark.macro.ExperimentalMetricApi
 import androidx.benchmark.macro.FrameTimingMetric
@@ -26,8 +25,8 @@ import androidx.benchmark.macro.TraceSectionMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -54,18 +53,18 @@ class FrameTimingBenchmark {
             iterations = DEFAULT_ITERATIONS,
             // [END_EXCLUDE]
             setupBlock = {
-                // Before starting to measure, navigate to the UI to be measured
-                val intent = Intent("$packageName.RECYCLER_VIEW_ACTIVITY")
-                startActivityAndWait(intent)
+                uiAutomator {
+                    // Before starting to measure, navigate to the UI to be measured
+                    startIntent(Intent("$packageName.RECYCLER_VIEW_ACTIVITY"))
+                }
             }
         ) {
-            val recycler = device.findObject(By.res(packageName, "recycler"))
-            // Set gesture margin to avoid triggering gesture navigation
-            // with input events from automation.
-            recycler.setGestureMargin(device.displayWidth / 5)
+            uiAutomator {
+                val recycler = onElement { className == "androidx.recyclerview.widget.RecyclerView" }
+                // Scroll down several times
+                repeat(3) { recycler.fling(Direction.DOWN) }
+            }
 
-            // Scroll down several times
-            repeat(3) { recycler.fling(Direction.DOWN) }
         }
     }
     // [END macrobenchmark_control_your_app]
@@ -96,27 +95,15 @@ class FrameTimingBenchmark {
             iterations = DEFAULT_ITERATIONS,
             // [END_EXCLUDE]
             setupBlock = {
-                // Before starting to measure, navigate to the UI to be measured
-                val intent = Intent("$packageName.COMPOSE_ACTIVITY")
-                startActivityAndWait(intent)
+                uiAutomator {
+                    // Before starting to measure, navigate to the UI to be measured
+                    startIntent(Intent("$packageName.COMPOSE_ACTIVITY"))
+                }
             }
         ) {
-            /**
-             * Compose does not have view IDs so we cannot directly access composables from UiAutomator.
-             * To access a composable we need to set:
-             * 1) Modifier.semantics { testTagsAsResourceId = true } once, high in the compose hierarchy
-             * 2) Add Modifier.testTag("someIdentifier") to all of the composables you want to access
-             *
-             * Once done that, we can access the composable using By.res("someIdentifier")
-             */
-            val column = device.findObject(By.res("myLazyColumn"))
-
-            // Set gesture margin to avoid triggering gesture navigation
-            // with input events from automation.
-            column.setGestureMargin(device.displayWidth / 5)
-
-            // Scroll down several times
-            repeat(1) { column.drag(Point(column.visibleCenter.x, column.visibleBounds.top)) }
+            uiAutomator {
+                onElement { isScrollable }.fling(Direction.DOWN)
+            }
         }
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NestedRecyclerFrameTimingBenchmarks.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NestedRecyclerFrameTimingBenchmarks.kt
index e8cb0ea0..c7176f2c 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NestedRecyclerFrameTimingBenchmarks.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NestedRecyclerFrameTimingBenchmarks.kt
@@ -16,22 +16,21 @@
 
 package com.example.macrobenchmark.benchmark.frames
 
-import androidx.benchmark.macro.*
+import androidx.benchmark.macro.FrameTimingMetric
+import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
-import androidx.test.uiautomator.Until
+import androidx.test.uiautomator.onElement
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
-import com.example.macrobenchmark.benchmark.util.waitAndFind
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import java.util.concurrent.TimeUnit
 
-@ExperimentalMetricApi
 @LargeTest
 @RunWith(AndroidJUnit4::class)
 class NestedRecyclerFrameTimingBenchmarks {
@@ -60,44 +59,24 @@ class NestedRecyclerFrameTimingBenchmarks {
         ) { measureScrollingNestedRecycler() }
     }
 
-    private fun MacrobenchmarkScope.navigateToNestedRvScreen(useRecyclerViewPool: Boolean) {
-        startActivityAndWait()
-
-        // navigate to the activity
-        val selector =
-            By.res(if (useRecyclerViewPool) "nestedRecyclerWithPoolsActivity" else "nestedRecyclerActivity")
-
-        val button = device.waitAndFind(selector)
-        button.click()
-
-        // wait until the activity is shown
-        device.wait(
-            Until.hasObject(By.clazz("$packageName.NestedRecyclerActivity")),
-            TimeUnit.SECONDS.toMillis(10)
-        )
+    private fun navigateToNestedRvScreen(useRecyclerViewPool: Boolean) {
+        uiAutomator {
+            startApp(TARGET_PACKAGE)
+            onElement { textAsString() == if (useRecyclerViewPool) "Nested RecyclerView with Pools" else "Nested RecyclerView" }.click()
+            waitForStableInActiveWindow()
+        }
     }
 
-    private fun MacrobenchmarkScope.measureScrollingNestedRecycler() {
-        val recycler = device.findObject(By.res(packageName, "recycler"))
-        // Set gesture margin to avoid triggering gesture navigation with input events from automation.
-        recycler.setGestureMargin(device.displayWidth / 5)
-
-        repeat(3) { index ->
-            val visibleNestedRecyclers =
-                recycler.findObjects(By.res(packageName, "row_recycler"))
-
-            // scroll the second recycler, because the first one may be shown just a pixel
-            val nestedRecyclerToScroll = visibleNestedRecyclers[1]
-
-            // Set gesture margin to avoid triggering gesture navigation with input events from automation.
-            nestedRecyclerToScroll.setGestureMargin(device.displayWidth / 3)
-
-            // swipe horizontally
-            nestedRecyclerToScroll.fling(Direction.RIGHT)
-            // scroll down twice and once up
-            recycler.swipe(if (index < 2) Direction.UP else Direction.DOWN, 0.5f)
-            // wait until the swipe is done
-            device.waitForIdle()
+    private fun measureScrollingNestedRecycler() {
+        uiAutomator {
+            onElement { className == "androidx.recyclerview.widget.RecyclerView" }.run {
+                repeat(3) { index ->
+                    onElement { className == "androidx.recyclerview.widget.RecyclerView" }.fling(
+                        Direction.RIGHT
+                    )
+                    swipe(if (index < 2) Direction.UP else Direction.DOWN, 0.5f)
+                }
+            }
         }
     }
 
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NonExportedActivityBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NonExportedActivityBenchmark.kt
index ccb1855a..8216053e 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NonExportedActivityBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/NonExportedActivityBenchmark.kt
@@ -23,16 +23,14 @@ import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
-import androidx.test.uiautomator.Until
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
-import junit.framework.TestCase.fail
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import java.util.concurrent.TimeUnit
 
 
 @LargeTest
@@ -60,37 +58,24 @@ class NonExportedActivityBenchmark {
             setupBlock = setupBenchmark()
         ) {
             // [START_EXCLUDE]
-            val recycler = device.findObject(By.res(packageName, "recycler"))
-
-            // Set gesture margin to avoid triggering gesture navigation
-            // with input events from automation.
-            recycler.setGestureMargin(device.displayWidth / 5)
-
-            // Fling the recycler several times
-            repeat(3) { recycler.fling(Direction.DOWN) }
+            uiAutomator {
+                onElement { className == "androidx.recyclerview.widget.RecyclerView" }.run {
+                    repeat(3) { fling(Direction.DOWN)}
+                }
+            }
             // [END_EXCLUDE]
         }
     }
 
     private fun setupBenchmark(): MacrobenchmarkScope.() -> Unit = {
-        // Before starting to measure, navigate to the UI to be measured
-        startActivityAndWait()
-
-        // click a button to launch the target activity.
-        // While we use button text  here to find the button, you could also use
-        // accessibility info or resourceId.
-        val selector = By.text("RecyclerView")
-        if (!device.wait(Until.hasObject(selector), 5_500)) {
-            fail("Could not find resource in time")
+        uiAutomator {
+            // Before starting to measure, navigate to the UI to be measured
+            startApp(TARGET_PACKAGE)
+            // click a button to launch the target activity.
+            onElement { textAsString() == "RecyclerView" }.click()
+            // wait until the activity is shown
+            waitForStableInActiveWindow()
         }
-        val launchRecyclerActivity = device.findObject(selector)
-        launchRecyclerActivity.click()
-
-        // wait until the activity is shown
-        device.wait(
-            Until.hasObject(By.clazz("$packageName.NonExportedRecyclerActivity")),
-            TimeUnit.SECONDS.toMillis(10)
-        )
     }
     // [END macrobenchmark_navigate_within_app]
 
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/TextInputFrameTimingBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/TextInputFrameTimingBenchmark.kt
index b4d3b467..421796eb 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/TextInputFrameTimingBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/frames/TextInputFrameTimingBenchmark.kt
@@ -23,7 +23,8 @@ import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -37,7 +38,7 @@ class TextInputFrameTimingBenchmark {
     val benchmarkRule = MacrobenchmarkRule()
 
     @Test
-    fun inputText() {
+    fun reinputText() {
         benchmarkRule.measureRepeated(
             packageName = TARGET_PACKAGE,
             metrics = listOf(FrameTimingMetric()),
@@ -47,15 +48,14 @@ class TextInputFrameTimingBenchmark {
             startupMode = StartupMode.WARM, // restarts activity each iteration
             iterations = DEFAULT_ITERATIONS,
             setupBlock = {
-                // Before starting to measure, navigate to the UI to be measured
-                val intent = Intent("$packageName.COMPOSE_ACTIVITY")
-                startActivityAndWait(intent)
+                uiAutomator {
+                    // Before starting to measure, navigate to the UI to be measured
+                    startIntent(Intent("$packageName.COMPOSE_ACTIVITY"))
+                }
             }
         ) {
-            val input = device.findObject(By.res("input"))
-            repeat(3) {
-                input.text = ""
-                input.text = "Benchmark input $it"
+            uiAutomator {
+                onElement { textAsString() == "Enter text here" }.text = "Benchmark input"
             }
         }
     }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/permissions/Permissions.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/permissions/Permissions.kt
deleted file mode 100644
index 7a374c43..00000000
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/permissions/Permissions.kt
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * 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
- *
- *     https://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 com.example.macrobenchmark.benchmark.permissions
-
-import android.Manifest.permission
-import android.os.Build.VERSION.SDK_INT
-import android.os.Build.VERSION_CODES.TIRAMISU
-import androidx.benchmark.macro.MacrobenchmarkScope
-import org.junit.Assert
-
-/**
- * Because the app under test runs in a different process from the one running the instrumentation test,
- * the permission has to be granted manually by either:
- *
- * - tapping the Allow button
- *    ```kotlin
- *    val obj = By.text("Allow")
- *    val dialog = device.wait(Until.findObject(obj), TIMEOUT)
- *    dialog?.let {
- *        it.click()
- *        device.wait(Until.gone(obj), 5_000)
- *    }
- *    ```
- * - or (preferred) executing the grant command on the target package.
- */
-fun MacrobenchmarkScope.allowNotifications() {
-    if (SDK_INT >= TIRAMISU) {
-        val command = "pm grant $packageName ${permission.POST_NOTIFICATIONS}"
-        val output = device.executeShellCommand(command)
-        Assert.assertEquals("", output)
-    }
-}
\ No newline at end of file
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/scroll/ScrollBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/scroll/ScrollBenchmark.kt
index cd7d690b..13a3eea6 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/scroll/ScrollBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/scroll/ScrollBenchmark.kt
@@ -24,13 +24,10 @@ import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.TraceSectionMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SdkSuppress
-import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
-import androidx.test.uiautomator.Until
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
-import junit.framework.TestCase
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -43,7 +40,6 @@ class ScrollBenchmark {
     val benchmarkRule = MacrobenchmarkRule()
 
     @Test
-    @SdkSuppress(minSdkVersion = 24)
     fun noCompilation() = scroll(CompilationMode.None())
 
     @Test
@@ -66,25 +62,21 @@ class ScrollBenchmark {
             iterations = DEFAULT_ITERATIONS,
             setupBlock = {
                 if (firstStart) {
-                    val intent = Intent("$packageName.SCROLL_VIEW_ACTIVITY")
-                    startActivityAndWait(intent)
+                    uiAutomator {
+                        startIntent(Intent("$packageName.SCROLL_VIEW_ACTIVITY"))
+                    }
                     firstStart = false
                 }
             }
         ) {
-            device.wait(Until.hasObject(By.scrollable(true)), 5_000)
-
-            val scrollableObject = device.findObject(By.scrollable(true))
-            if (scrollableObject == null) {
-                TestCase.fail("No scrollable view found in hierarchy")
-            }
-            scrollableObject.setGestureMargin(device.displayWidth / 10)
-            scrollableObject?.apply {
-                repeat(2) {
-                    fling(Direction.DOWN)
-                }
-                repeat(2) {
-                    fling(Direction.UP)
+            uiAutomator {
+                onElement { isScrollable }.run {
+                    repeat(2) {
+                        fling(Direction.DOWN)
+                    }
+                    repeat(2) {
+                        fling(Direction.UP)
+                    }
                 }
             }
         }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/FullyDrawnStartupBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/FullyDrawnStartupBenchmark.kt
index 576f577a..db823c11 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/FullyDrawnStartupBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/FullyDrawnStartupBenchmark.kt
@@ -22,8 +22,8 @@ import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
-import androidx.test.uiautomator.By
-import androidx.test.uiautomator.Until
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -43,12 +43,11 @@ class FullyDrawnStartupBenchmark {
         startupMode = StartupMode.COLD,
         iterations = DEFAULT_ITERATIONS,
     ) {
-        val intent = Intent("$TARGET_PACKAGE.FULLY_DRAWN_STARTUP_ACTIVITY")
+        uiAutomator {
+            startIntent(Intent("$TARGET_PACKAGE.FULLY_DRAWN_STARTUP_ACTIVITY"))
 
-        // Waits for first rendered frame
-        startActivityAndWait(intent)
-
-        // Waits for an element that corresponds to fully drawn state
-        device.wait(Until.hasObject(By.text("Fully Drawn")), 10_000)
+            // Waits for an element that corresponds to fully drawn state
+            onElement { textAsString() == "Fully Drawn" }
+        }
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SampleStartupBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SampleStartupBenchmark.kt
index 2f299059..dd3fae26 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SampleStartupBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SampleStartupBenchmark.kt
@@ -20,6 +20,7 @@ import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.LargeTest
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -44,7 +45,7 @@ class SampleStartupBenchmark {
         }
     ) {
         // starts default launch activity
-        startActivityAndWait()
+        uiAutomator { startApp(TARGET_PACKAGE) }
     }
 }
 // [END macrobenchmark_startup]
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SmallListStartupBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SmallListStartupBenchmark.kt
index 47cf897a..f64b6741 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SmallListStartupBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/SmallListStartupBenchmark.kt
@@ -21,6 +21,7 @@ import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
 import androidx.test.filters.LargeTest
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -49,10 +50,14 @@ class SmallListStartupBenchmark(
             pressHome()
         }
     ) {
-        val intent = Intent("$packageName.RECYCLER_VIEW_ACTIVITY")
-        intent.putExtra("ITEM_COUNT", LIST_ITEMS)
-
-        startActivityAndWait(intent)
+        uiAutomator {
+            startIntent(
+                Intent("$packageName.RECYCLER_VIEW_ACTIVITY").putExtra(
+                    "ITEM_COUNT",
+                    LIST_ITEMS
+                )
+            )
+        }
     }
 
     @Test
@@ -66,10 +71,12 @@ class SmallListStartupBenchmark(
             pressHome()
         }
     ) {
-        val intent = Intent("$packageName.COMPOSE_ACTIVITY")
-        intent.putExtra("ITEM_COUNT", LIST_ITEMS)
-
-        startActivityAndWait(intent)
+        uiAutomator {
+            startIntent(
+                Intent("$packageName.COMPOSE_ACTIVITY")
+                    .putExtra("ITEM_COUNT", LIST_ITEMS)
+            )
+        }
     }
 
     companion object {
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/StartupBenchmark.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/StartupBenchmark.kt
index 0da8cfb8..ee0f06a3 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/StartupBenchmark.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/startup/StartupBenchmark.kt
@@ -21,8 +21,9 @@ import androidx.benchmark.macro.CompilationMode
 import androidx.benchmark.macro.StartupMode
 import androidx.benchmark.macro.StartupTimingMetric
 import androidx.benchmark.macro.junit4.MacrobenchmarkRule
-import androidx.test.filters.SdkSuppress
 import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
+import androidx.test.uiautomator.textAsString
+import androidx.test.uiautomator.uiAutomator
 import com.example.macrobenchmark.benchmark.util.DEFAULT_ITERATIONS
 import com.example.macrobenchmark.benchmark.util.TARGET_PACKAGE
 import org.junit.Rule
@@ -60,11 +61,9 @@ abstract class AbstractStartupBenchmark(private val startupMode: StartupMode) {
     val benchmarkRule = MacrobenchmarkRule()
 
     @Test
-    @SdkSuppress(minSdkVersion = 24)
     fun startupNoCompilation() = startup(CompilationMode.None())
 
     @Test
-    @SdkSuppress(minSdkVersion = 24)
     fun startupPartialCompilation() = startup(
         CompilationMode.Partial(
             baselineProfileMode = BaselineProfileMode.Disable,
@@ -73,7 +72,6 @@ abstract class AbstractStartupBenchmark(private val startupMode: StartupMode) {
     )
 
     @Test
-    @SdkSuppress(minSdkVersion = 24)
     fun startupPartialWithBaselineProfiles() =
         startup(CompilationMode.Partial(baselineProfileMode = BaselineProfileMode.Require))
 
@@ -90,6 +88,9 @@ abstract class AbstractStartupBenchmark(private val startupMode: StartupMode) {
             pressHome()
         }
     ) {
-        startActivityAndWait()
+        uiAutomator {
+            startApp(TARGET_PACKAGE)
+            onElement { textAsString() == "Fully Drawn" && isVisibleToUser }
+        }
     }
 }
\ No newline at end of file
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/util/Utils.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/util/Utils.kt
index dbee5033..eef29c59 100644
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/util/Utils.kt
+++ b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/benchmark/util/Utils.kt
@@ -16,69 +16,5 @@
 
 package com.example.macrobenchmark.benchmark.util
 
-import android.util.Log
-import androidx.test.uiautomator.BySelector
-import androidx.test.uiautomator.SearchCondition
-import androidx.test.uiautomator.UiDevice
-import androidx.test.uiautomator.UiObject2
-import androidx.test.uiautomator.Until
-import java.io.ByteArrayOutputStream
-
 const val TARGET_PACKAGE = "com.example.macrobenchmark.target"
-const val DEFAULT_ITERATIONS = 10
-
-/**
- * Finds an element by [selector].
- * If not found, fails with [AssertionError] and dumps the window hierarchy.
- */
-fun UiDevice.findOrFail(
-    selector: BySelector,
-    message: String? = null,
-): UiObject2 {
-    val element = findObject(selector)
-    if (element == null) {
-        val hierarchy = getWindowHierarchy()
-        Log.d("Benchmark", hierarchy)
-        throw AssertionError((message ?: "Object not on screen ($selector)") + "\n$hierarchy")
-    }
-    return element
-}
-
-/**
- * Waits until a [searchCondition] evaluates to true.
- * If not found within [timeout], fails with [AssertionError] and dumps the window hierarchy.
- * For example, wait [Until.hasObject] to wait until an element is present on screen.
- */
-fun UiDevice.waitOrFail(
-    searchCondition: SearchCondition<Boolean>,
-    timeout: Long,
-    message: String? = null,
-) {
-    if (!wait(searchCondition, timeout)) {
-        val hierarchy = getWindowHierarchy()
-        Log.d("Benchmark", hierarchy)
-        throw AssertionError((message ?: "Object not on screen") + "\n$hierarchy")
-    }
-}
-
-/**
- * Combines waiting for an element and returning it.
- * If an object is not present, it throws [AssertionError] and dumps the window hierarchy.
- */
-fun UiDevice.waitAndFind(
-    selector: BySelector,
-    timeout: Long = 5_000,
-    message: String? = null,
-): UiObject2 {
-    waitOrFail(Until.hasObject(selector), timeout, message)
-    return findOrFail(selector, message)
-}
-
-/**
- * Simplifies dumping window hierarchy
- */
-fun UiDevice.getWindowHierarchy(): String {
-    val output = ByteArrayOutputStream()
-    dumpWindowHierarchy(output)
-    return output.toString()
-}
+const val DEFAULT_ITERATIONS = 1