diff --git a/ui/ui-material/src/androidTest/java/androidx/ui/material/DrawerTest.kt b/ui/ui-material/src/androidTest/java/androidx/ui/material/DrawerTest.kt
index 61a1c8a..64f08a7 100644
--- a/ui/ui-material/src/androidTest/java/androidx/ui/material/DrawerTest.kt
+++ b/ui/ui-material/src/androidTest/java/androidx/ui/material/DrawerTest.kt
@@ -174,21 +174,22 @@
         findByTag("Drawer")
             .doClick()
 
-        Truth.assertThat(drawerClicks).isEqualTo(0)
-        Truth.assertThat(bodyClicks).isEqualTo(1)
+        composeTestRule.runOnIdleCompose {
+            Truth.assertThat(drawerClicks).isEqualTo(0)
+            Truth.assertThat(bodyClicks).isEqualTo(1)
+        }
 
         composeTestRule.runOnUiThread {
             drawerState.state = DrawerState.Opened
         }
-        // TODO: we aren't correctly waiting for recompositions after clicking, so we need to wait
-        // again
-        Thread.sleep(100L)
 
         findByTag("Drawer")
             .doClick()
 
-        Truth.assertThat(drawerClicks).isEqualTo(1)
-        Truth.assertThat(bodyClicks).isEqualTo(1)
+        composeTestRule.runOnIdleCompose {
+            Truth.assertThat(drawerClicks).isEqualTo(1)
+            Truth.assertThat(bodyClicks).isEqualTo(1)
+        }
     }
 
     @Test