diff --git a/ui/integration-tests/test/src/main/java/androidx/ui/test/cases/RectsInColumnSharedModelTestCase.kt b/ui/integration-tests/test/src/main/java/androidx/ui/test/cases/RectsInColumnSharedModelTestCase.kt
index 0dcf66c..c7ddc466 100644
--- a/ui/integration-tests/test/src/main/java/androidx/ui/test/cases/RectsInColumnSharedModelTestCase.kt
+++ b/ui/integration-tests/test/src/main/java/androidx/ui/test/cases/RectsInColumnSharedModelTestCase.kt
@@ -16,11 +16,9 @@
 
 package androidx.ui.test.cases
 
-import android.app.Activity
-import androidx.compose.FrameManager
+import androidx.compose.Composable
 import androidx.compose.Model
 import androidx.ui.core.dp
-import androidx.ui.core.setContent
 import androidx.ui.foundation.ColoredRect
 import androidx.ui.graphics.Color
 import androidx.ui.layout.Column
@@ -39,13 +37,13 @@
  * that the whole loop has to be re-run when model changes.
  */
 class RectsInColumnSharedModelTestCase(
-    activity: Activity,
     private val amountOfRectangles: Int
-) : ComposeTestCase(activity), ToggleableTestCase {
+) : ComposeTestCase, ToggleableTestCase {
 
     private val model = RectanglesInColumnTestCaseColorModel(Color.Black)
 
-    override fun setComposeContent(activity: Activity) = activity.setContent {
+    @Composable
+    override fun emitContent() {
         MaterialTheme {
             Column {
                 repeat(amountOfRectangles) { i ->
@@ -57,7 +55,7 @@
                 }
             }
         }
-    }!!
+    }
 
     override fun toggleState() {
         if (model.color == Color.Magenta) {
@@ -65,6 +63,5 @@
         } else {
             model.color = Color.Magenta
         }
-        FrameManager.nextFrame()
     }
 }
\ No newline at end of file