From 6291dd79e64e7a85253a9414d877e191133dda06 Mon Sep 17 00:00:00 2001
From: Jacky Wang <jiannan@google.com>
Date: Mon, 23 Dec 2024 15:31:45 +0800
Subject: [PATCH] Fix flaky LifecycleAwareAsyncTaskTest

Fix: 385137513
Flag: TEST_ONLY
Test: atest SettingsRoboTests:LifecycleAwareAsyncTaskTest --rerun-until-failure 500
Change-Id: I269d2fa1c63e0b805d7d1b0710919f60c2096ec0
---
 .../batteryusage/LifecycleAwareAsyncTask.kt   | 16 ++++++-----
 .../LifecycleAwareAsyncTaskTest.kt            | 27 +++++++++----------
 2 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt b/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt
index a715cb7f363..6b5308ccad7 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt
+++ b/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTask.kt
@@ -18,6 +18,7 @@ package com.android.settings.fuelgauge.batteryusage
 
 import android.os.AsyncTask
 import androidx.annotation.CallSuper
+import androidx.annotation.OpenForTesting
 import androidx.lifecycle.DefaultLifecycleObserver
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
@@ -50,12 +51,15 @@ abstract class LifecycleAwareAsyncTask<Result>(private val lifecycle: Lifecycle?
     fun start() {
         execute() // expects main thread
         val lifecycle = lifecycle ?: return
-        mainExecutor.execute {
-            // Status is updated to FINISHED if onPoseExecute happened before. And task is cancelled
-            // if lifecycle is stopped.
-            if (status == Status.RUNNING && !isCancelled) {
-                lifecycle.addObserver(this) // requires main thread
-            }
+        mainExecutor.execute { maybeAddObserver(lifecycle) }
+    }
+
+    @OpenForTesting
+    open fun maybeAddObserver(lifecycle: Lifecycle) {
+        // Status is updated to FINISHED if onPoseExecute happened before. And task is cancelled
+        // if lifecycle is stopped.
+        if (status == Status.RUNNING && !isCancelled) {
+            lifecycle.addObserver(this) // requires main thread
         }
     }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt
index fde45b7232d..843a55fc501 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batteryusage/LifecycleAwareAsyncTaskTest.kt
@@ -22,12 +22,13 @@ import androidx.lifecycle.LifecycleOwner
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit.MILLISECONDS
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.verify
-import java.util.concurrent.CountDownLatch
 
 @RunWith(AndroidJUnit4::class)
 class LifecycleAwareAsyncTaskTest {
@@ -85,6 +86,7 @@ class LifecycleAwareAsyncTaskTest {
 
     @Test
     fun onPostExecute_addObserver() {
+        val countDownLatch = CountDownLatch(2)
         val observers = mutableListOf<LifecycleObserver>()
         val lifecycle =
             object : Lifecycle() {
@@ -97,32 +99,27 @@ class LifecycleAwareAsyncTaskTest {
 
                 override fun removeObserver(observer: LifecycleObserver) {
                     observers.remove(observer)
+                    countDownLatch.countDown()
                 }
             }
         val asyncTask =
             object : LifecycleAwareAsyncTask<Void?>(lifecycle) {
                 override fun doInBackground(vararg params: Void) = null
+
+                override fun maybeAddObserver(lifecycle: Lifecycle) {
+                    super.maybeAddObserver(lifecycle)
+                    countDownLatch.countDown()
+                }
             }
 
         Thread { asyncTask.start() }.start()
-        idleAsyncTaskExecutor()
-        instrumentation.waitForIdleSync()
+        do {
+            instrumentation.waitForIdleSync()
+        } while (!countDownLatch.await(100, MILLISECONDS))
 
         assertThat(observers).isEmpty()
     }
 
-    private fun idleAsyncTaskExecutor() {
-        val taskCountDownLatch = CountDownLatch(1)
-        object : LifecycleAwareAsyncTask<Void?>(null) {
-                override fun doInBackground(vararg params: Void): Void? {
-                    taskCountDownLatch.countDown()
-                    return null
-                }
-            }
-            .start()
-        taskCountDownLatch.await()
-    }
-
     @Test
     fun onStop_addObserver() {
         val executorBlocker = CountDownLatch(1)