diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/BaselineProfileGeneratorScaffold.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/BaselineProfileGeneratorScaffold.kt
deleted file mode 100644
index 53b3562..0000000
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/BaselineProfileGeneratorScaffold.kt
+++ /dev/null
@@ -1,52 +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.benchmark.macro.junit4.BaselineProfileRule
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import org.junit.Rule
-import org.junit.Test
-import org.junit.runner.RunWith
-
-/**
- * A scaffold for creating a baseline profile user journey. Implementing classes can
- * start generating a profile directly by implementing [MacrobenchmarkScope.profileBlock].
- */
-@RunWith(AndroidJUnit4::class)
-abstract class BaselineProfileGeneratorScaffold {
-
-    @get:Rule
-    val rule = BaselineProfileRule()
-    
-    /**
-     * Generate a baseline profile in this function.
-     */
-    abstract fun MacrobenchmarkScope.profileBlock()
-
-    @Test
-    fun profileGenerator() {
-        rule.collect(
-            packageName = TARGET_PACKAGE,
-            maxIterations = 15,
-            stableIterations = 3
-        ) {
-            profileBlock()
-        }
-    }
-
-}
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 c6f2d3a..7264400 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
@@ -17,29 +17,41 @@
 package com.example.macrobenchmark.baselineprofile
 
 import android.content.Intent
-import androidx.benchmark.macro.MacrobenchmarkScope
-import androidx.test.ext.junit.runners.AndroidJUnit4
+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 org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
 import org.junit.runner.RunWith
 
 @Ignore // TODO causing stale object excpetion on CI .. why?
-@RunWith(AndroidJUnit4::class)
-class ComposeActivityBaselineProfileGenerator : BaselineProfileGeneratorScaffold() {
+@RunWith(AndroidJUnit4ClassRunner::class)
+class ComposeActivityBaselineProfileGenerator {
 
-    override fun MacrobenchmarkScope.profileBlock() {
-        // Start into the Compose Activity
-        startActivityAndWait(Intent("$TARGET_PACKAGE.COMPOSE_ACTIVITY"))
+    @get:Rule
+    val rule = BaselineProfileRule()
 
-        // Scrolling through the Compose journey
-        device.waitAndFind(By.res("myLazyColumn")).also {
-            it.setGestureMargin(device.displayWidth / 10)
-            it.fling(Direction.DOWN)
-        }
+    @Test
+    fun generate() {
+        rule.collect(
+            packageName = TARGET_PACKAGE,
+            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)
+            }
 
-        device.findOrFail(By.res("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 3eaf92d..eff0a37 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
@@ -17,16 +17,31 @@
 package com.example.macrobenchmark.baselineprofile
 
 import android.content.Intent
-import androidx.benchmark.macro.MacrobenchmarkScope
+import androidx.benchmark.macro.junit4.BaselineProfileRule
+import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
 import androidx.test.uiautomator.By
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
 
-class LoginBaselineProfileGenerator : BaselineProfileGeneratorScaffold() {
+@RunWith(AndroidJUnit4ClassRunner::class)
+class LoginBaselineProfileGenerator {
 
-    override fun MacrobenchmarkScope.profileBlock() {
-        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()
+    @get:Rule
+    val rule = BaselineProfileRule()
+
+    @Test
+    fun generate() {
+        rule.collect(
+            packageName = TARGET_PACKAGE,
+            maxIterations = 15,
+            stableIterations = 3
+        ) {
+            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()
+        }
     }
 }
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 16da674..a5ab1b8 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
@@ -17,26 +17,38 @@
 package com.example.macrobenchmark.baselineprofile
 
 import android.content.Intent
-import androidx.benchmark.macro.MacrobenchmarkScope
-import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.benchmark.macro.junit4.BaselineProfileRule
+import androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Direction
 import org.junit.Ignore
+import org.junit.Rule
+import org.junit.Test
 import org.junit.runner.RunWith
 
 @Ignore // TODO flinging not working, ignore for now to test the pipeline
-@RunWith(AndroidJUnit4::class)
-class RecyclerViewActivityBaselineProfileGenerator : BaselineProfileGeneratorScaffold() {
+@RunWith(AndroidJUnit4ClassRunner::class)
+class RecyclerViewActivityBaselineProfileGenerator {
 
-    override fun MacrobenchmarkScope.profileBlock() {
-        // Start into the RecyclerViewActivity
-        startActivityAndWait(Intent("$TARGET_PACKAGE.RECYCLER_VIEW_ACTIVITY"))
+    @get:Rule
+    val rule = BaselineProfileRule()
 
-        // Scrolling RecyclerView journey
-        device.findObject(By.res(packageName, "recycler")).also {
-            it.setGestureMargin(device.displayWidth / 10)
-            it.fling(Direction.DOWN)
-            it.fling(Direction.UP)
+    @Test
+    fun generate() {
+        rule.collect(
+            packageName = TARGET_PACKAGE,
+            maxIterations = 15,
+            stableIterations = 3
+        ) {
+            // Start into the RecyclerViewActivity
+            startActivityAndWait(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)
+            }
         }
     }
 }
diff --git a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupOnlyBaselineProfileGenerator.kt b/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupOnlyBaselineProfileGenerator.kt
deleted file mode 100644
index 8e1f40f..0000000
--- a/MacrobenchmarkSample/macrobenchmark/src/main/kotlin/com/example/macrobenchmark/baselineprofile/StartupOnlyBaselineProfileGenerator.kt
+++ /dev/null
@@ -1,29 +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
-
-/**
- * This baseline profile generator creates a baseline profile only for the app startup.
- */
-class StartupOnlyBaselineProfileGenerator: BaselineProfileGeneratorScaffold() {
-
-    override fun MacrobenchmarkScope.profileBlock() {
-        startActivityAndWait()
-    }
-}