Fix ImageViewAssistantTest timeout

Instead of relying on a timeout, use the future to deterministically
detect the completion of the component initialization.

Fixes: 192271062
Test: ImageViewAssistantTest
Change-Id: I209d887fa3cde183a60cc7f9732c4f17620bd31b

2 files changed, 11 insertions, 9 deletions

diff --git a/designer/src/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistant.kt b/designer/src/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistant.kt
index d8cb4e4b152..43dfa162d38 100644
--- a/designer/src/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistant.kt
+++ b/designer/src/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistant.kt
@@ -37,6 +37,7 @@ import com.intellij.util.concurrency.EdtExecutorService
 import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.JBUI.Borders
 import com.intellij.util.ui.JBUI.scale
+import org.jetbrains.annotations.TestOnly
 import java.awt.BorderLayout
 import java.util.EnumSet
 import java.util.concurrent.CompletableFuture
@@ -68,6 +69,12 @@ class ImageViewAssistant(
 
   private val itemNameLabel = assistantLabel(getSampleItemDisplayName(originalValue))
 
+  /**
+   * [CompletableFuture] used to verify that the load of the sample data resources is complete.
+   */
+  @TestOnly
+  val sampleDataLoaded: CompletableFuture<List<SampleDataResourceItem>>
+
   private var itemDisplayName: String?
     get() = itemNameLabel.text
     set(itemName) {
@@ -113,9 +120,9 @@ class ImageViewAssistant(
     updateUIState()
 
     // Get SampleData drawables in background thread, then, update the widget on the EDT
-    CompletableFuture.supplyAsync(Supplier {
+    sampleDataLoaded = CompletableFuture.supplyAsync({
       ResourceRepositoryManager.getAppResources(nlComponent.model.facet).getSampleDataOfType(IMAGE).toList()
-    }, AppExecutorUtil.getAppExecutorService()).whenCompleteAsync(BiConsumer { sampleDataItems, _ ->
+    }, AppExecutorUtil.getAppExecutorService()).whenCompleteAsync({ sampleDataItems, _ ->
       populateWidget(sampleDataItems)
     }, EdtExecutorService.getScheduledExecutorInstance())
   }
diff --git a/designer/testSrc/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistantTest.kt b/designer/testSrc/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistantTest.kt
index 649d6b75b8b..38f30123f11 100644
--- a/designer/testSrc/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistantTest.kt
+++ b/designer/testSrc/com/android/tools/idea/uibuilder/handlers/assistant/ImageViewAssistantTest.kt
@@ -22,7 +22,6 @@ import com.android.tools.idea.testing.AndroidProjectRule
 import com.android.tools.idea.uibuilder.assistant.ComponentAssistantFactory
 import com.android.tools.idea.uibuilder.handlers.ImageViewHandler
 import com.google.common.truth.Truth.assertThat
-import com.intellij.util.WaitFor
 import com.intellij.util.ui.UIUtil
 import org.jetbrains.android.facet.AndroidFacet
 import org.junit.Rule
@@ -60,12 +59,8 @@ class ImageViewAssistantTest {
     val assistant = ImageViewAssistant(ComponentAssistantFactory.Context(nlComponent) {}, imageHandler)
     val content = assistant.component.content
 
-    with(object : WaitFor(200) {
-      override fun condition(): Boolean {
-        // Enabled ComboBox means it finished loading the sample data into the ui
-        return UIUtil.findComponentOfType(content, JComboBox::class.java)?.isEnabled == true
-      }
-    }) { assertTrue(isConditionRealized) }
+    assistant.sampleDataLoaded.get()
+    assertTrue(UIUtil.findComponentOfType(content, JComboBox::class.java)?.isEnabled == true)
 
     val useAllCheckBox = UIUtil.findComponentOfType(content, JCheckBox::class.java)!!
     val sampleSetComboBox = UIUtil.findComponentOfType(content, JComboBox::class.java)!!