From 3a1af7208b19fea105d141eca8fcc658c2bbe5c9 Mon Sep 17 00:00:00 2001
From: Mady Mellor <madym@google.com>
Date: Thu, 6 Jul 2023 16:11:47 -0700
Subject: [PATCH] Wait for bubbles to be loaded in BubbleDataRepositoryTest

I made some changes ag/23755383 where we might write to XML again
when loading bubbles to clear out anything that is no longer valid.

Testing this is causing some issues with another test that validates
XML is read / written correctly. Basically, the load is async so
it might write over data being used in another test. The fix
is to wait until the load is finished in BubbleDataRepositoryTest.

Additionally uses a non-mocked executor so that the callback actually
runs, before it wasn't so this test wasn't really testing properly.

Test: atest BubbleDataRepositoryTest BubblePersistentRepositoryTest --iterations 10
Bug: 288984801
Bug: 288730033
Change-Id: I461287ec40ca0a4e01185a108a23c4f25b401325
---
 .../shell/bubbles/BubbleDataRepositoryTest.kt | 53 ++++++++++++-------
 1 file changed, 35 insertions(+), 18 deletions(-)

diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataRepositoryTest.kt
index 6d9d62d3ad925..e525ab35d0ddd 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataRepositoryTest.kt
@@ -19,17 +19,21 @@ package com.android.wm.shell.bubbles
 import android.app.ActivityTaskManager
 import android.content.pm.LauncherApps
 import android.content.pm.ShortcutInfo
+import android.os.Handler
+import android.os.Looper
 import android.util.SparseArray
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.bubbles.storage.BubbleEntity
 import com.android.wm.shell.bubbles.storage.BubblePersistentRepository
-import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.common.HandlerExecutor
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.spy
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 
 class BubbleDataRepositoryTest : ShellTestCase() {
 
@@ -118,7 +122,8 @@ class BubbleDataRepositoryTest : ShellTestCase() {
         )
     )
 
-    private val mainExecutor = mock(ShellExecutor::class.java)
+    private val testHandler = Handler(Looper.getMainLooper())
+    private val mainExecutor = HandlerExecutor(testHandler)
     private val launcherApps = mock(LauncherApps::class.java)
 
     private val persistedBubbles = SparseArray<List<BubbleEntity>>()
@@ -145,32 +150,44 @@ class BubbleDataRepositoryTest : ShellTestCase() {
 
     @Test
     fun testLoadBubbles_invalidParent() {
+        val latch = CountDownLatch(1)
         val activeUserIds = listOf(10, 1, 12) // Missing user 0 in persistedBubbles
+
         dataRepository.loadBubbles(1, activeUserIds) {
-            // Verify that user 0 has been removed from the persisted list
-            val entitiesByUser = persistentRepository.readFromDisk()
-            assertThat(entitiesByUser.get(0)).isNull()
+            latch.countDown()
         }
+        latch.await(500, TimeUnit.MILLISECONDS)
+        assertThat(latch.count).isEqualTo(0)
+
+        // Verify that user 0 has been removed from the persisted list
+        val entitiesByUser = persistentRepository.readFromDisk()
+        assertThat(entitiesByUser.get(0)).isNull()
     }
 
     @Test
     fun testLoadBubbles_invalidChild() {
+        val latch = CountDownLatch(1)
         val activeUserIds = listOf(0, 10, 1) // Missing user 1's child user 12
-        dataRepository.loadBubbles(1, activeUserIds) {
-            // Build a list to compare against
-            val user1BubblesWithoutUser12 = mutableListOf<Bubble>()
-            val user1EntitiesWithoutUser12 = mutableListOf<BubbleEntity>()
-            for (entity in user1BubbleEntities) {
-                if (entity.userId != 12) {
-                    user1BubblesWithoutUser12.add(entity.toBubble())
-                    user1EntitiesWithoutUser12.add(entity)
-                }
-            }
 
-            // Verify that user 12 has been removed from the persisted list
-            val entitiesByUser = persistentRepository.readFromDisk()
-            assertThat(entitiesByUser.get(1)).isEqualTo(user1EntitiesWithoutUser12)
+        // Build a list to compare against (remove all user 12 bubbles)
+        val user1BubblesWithoutUser12 = mutableListOf<Bubble>()
+        val user1EntitiesWithoutUser12 = mutableListOf<BubbleEntity>()
+        for (entity in user1BubbleEntities) {
+            if (entity.userId != 12) {
+                user1BubblesWithoutUser12.add(entity.toBubble())
+                user1EntitiesWithoutUser12.add(entity)
+            }
         }
+
+        dataRepository.loadBubbles(1, activeUserIds) {
+            latch.countDown()
+        }
+        latch.await(500, TimeUnit.MILLISECONDS)
+        assertThat(latch.count).isEqualTo(0)
+
+        // Verify that user 12 has been removed from the persisted list
+        val entitiesByUser = persistentRepository.readFromDisk()
+        assertThat(entitiesByUser.get(1)).isEqualTo(user1EntitiesWithoutUser12)
     }
 
     private fun BubbleEntity.toBubble(): Bubble {