From 37182de670b36184cfa5f989ebe20c9637997d20 Mon Sep 17 00:00:00 2001
From: Sunny Goyal <sunnygoyal@google.com>
Date: Wed, 25 Sep 2024 10:49:47 -0700
Subject: [PATCH] Fixing flakiness in robo tests

Bug: 352288591
Flag: EXEMPT test fix
Test: Presubmit
Change-Id: I15dd91815f7ae4d3cf3c9e1fd4c2b5e07052ab3f
---
 .../util/MainThreadInitializedObject.java     | 36 ++++++++++++++++---
 .../launcher3/RoboObjectInitializer.kt        | 35 ------------------
 .../launcher3/util/LauncherModelHelper.java   |  4 +--
 3 files changed, 34 insertions(+), 41 deletions(-)
 delete mode 100644 tests/multivalentTests/src/com/android/launcher3/RoboObjectInitializer.kt

diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index e12ccbcf70..a7d5c130f3 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -108,6 +108,25 @@ public class MainThreadInitializedObject<T extends SafeCloseable> {
          */
         <T extends SafeCloseable> T getObject(MainThreadInitializedObject<T> object);
 
+
+        /**
+         * Put a value into cache, can be used to put mocked MainThreadInitializedObject
+         * instances.
+         */
+        <T extends SafeCloseable> void putObject(MainThreadInitializedObject<T> object, T value);
+
+        /**
+         * Returns whether this context should cleanup all objects when its destroyed or leave it
+         * to the GC.
+         * These objects can have listeners attached to the system server and mey not be able to get
+         * GCed themselves when running on a device.
+         * Some environments like Robolectric tear down the whole system at the end of the test,
+         * so manual cleanup may not be required.
+         */
+        default boolean shouldCleanUpOnDestroy() {
+            return true;
+        }
+
         @UiThread
         default <T extends SafeCloseable> T createObject(MainThreadInitializedObject<T> object) {
             return object.mProvider.get((Context) this);
@@ -138,7 +157,19 @@ public class MainThreadInitializedObject<T extends SafeCloseable> {
             return this;
         }
 
+        @Override
+        public boolean shouldCleanUpOnDestroy() {
+            return (getBaseContext().getApplicationContext() instanceof SandboxApplication sa)
+                    ? sa.shouldCleanUpOnDestroy() : true;
+        }
+
         public void onDestroy() {
+            if (shouldCleanUpOnDestroy()) {
+                cleanUpObjects();
+            }
+        }
+
+        protected void cleanUpObjects() {
             getAppComponent().getDaggerSingletonTracker().close();
             synchronized (mDestroyLock) {
                 // Destroy in reverse order
@@ -174,10 +205,7 @@ public class MainThreadInitializedObject<T extends SafeCloseable> {
             }
         }
 
-        /**
-         * Put a value into mObjectMap, can be used to put mocked MainThreadInitializedObject
-         * instances into SandboxContext.
-         */
+        @Override
         public <T extends SafeCloseable> void putObject(
                 MainThreadInitializedObject<T> object, T value) {
             mObjectMap.put(object, value);
diff --git a/tests/multivalentTests/src/com/android/launcher3/RoboObjectInitializer.kt b/tests/multivalentTests/src/com/android/launcher3/RoboObjectInitializer.kt
deleted file mode 100644
index c5f9f86859..0000000000
--- a/tests/multivalentTests/src/com/android/launcher3/RoboObjectInitializer.kt
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.launcher3
-
-import com.android.launcher3.util.MainThreadInitializedObject
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxApplication
-import com.android.launcher3.util.SafeCloseable
-
-/**
- * Initializes [MainThreadInitializedObject] instances for Robolectric tests.
- *
- * Unlike instrumentation tests, Robolectric creates a new application instance for each test, which
- * could cause the various static objects defined in [MainThreadInitializedObject] to leak. Thus, a
- * [SandboxApplication] for Robolectric tests can implement this interface to limit the lifecycle of
- * these objects to a single test.
- */
-interface RoboObjectInitializer {
-
-    /** Overrides an object with [type] to [value]. */
-    fun <T : SafeCloseable> initializeObject(type: MainThreadInitializedObject<T>, value: T)
-}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
index 2d53e29bd6..1baee03ddf 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -283,11 +283,11 @@ public class LauncherModelHelper {
         }
 
         @Override
-        public void onDestroy() {
+        protected void cleanUpObjects() {
             if (deleteContents(mDbDir)) {
                 mDbDir.delete();
             }
-            super.onDestroy();
+            super.cleanUpObjects();
         }
 
         @Override