Migrating all model tests to Instrumentation tests

Bug: 196825541
Test: Presubmit
Change-Id: Iebd46eb41eb46c187d569197f4b97b4fddc0f6f7
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 702b73a..5080824 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -50,7 +50,7 @@
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
-public class LauncherAppState {
+public class LauncherAppState implements SafeCloseable {
 
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
     private static final String KEY_ICON_STATE = "pref_icon_shape_path";
@@ -158,7 +158,8 @@
     /**
      * Call from Application.onTerminate(), which is not guaranteed to ever be called.
      */
-    public void onTerminate() {
+    @Override
+    public void close() {
         mModel.destroy();
         mContext.getSystemService(LauncherApps.class).unregisterCallback(mModel);
         CustomWidgetManager.INSTANCE.get(mContext).setWidgetRefreshCallback(null);
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 7b6a5bf..9ebec0a 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -96,9 +96,10 @@
     // our monitoring of the package manager provides all updates and we never
     // need to do a requery. This is only ever touched from the loader thread.
     private boolean mModelLoaded;
+    private boolean mModelDestroyed = false;
     public boolean isModelLoaded() {
         synchronized (mLock) {
-            return mModelLoaded && mLoaderTask == null;
+            return mModelLoaded && mLoaderTask == null && !mModelDestroyed;
         }
     }
 
@@ -245,6 +246,7 @@
      * Called when the model is destroyed
      */
     public void destroy() {
+        mModelDestroyed = true;
         MODEL_EXECUTOR.execute(mModelDelegate::destroy);
     }
 
@@ -557,6 +559,9 @@
     }
 
     public void enqueueModelUpdateTask(ModelUpdateTask task) {
+        if (mModelDestroyed) {
+            return;
+        }
         task.init(mApp, this, mBgDataModel, mBgAllAppsList, MAIN_EXECUTOR);
         MODEL_EXECUTOR.execute(task);
     }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index a96de31..1c8954d 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -86,7 +86,7 @@
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
-import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
 import com.android.launcher3.widget.BaseLauncherAppWidgetHostView;
@@ -97,13 +97,10 @@
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 /**
@@ -122,22 +119,16 @@
      * Context used just for preview. It also provides a few objects (e.g. UserCache) just for
      * preview purposes.
      */
-    public static class PreviewContext extends ContextWrapper {
-
-        private final Set<MainThreadInitializedObject> mAllowedObjects = new HashSet<>(
-                Arrays.asList(UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
-                        LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                        CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE));
+    public static class PreviewContext extends SandboxContext {
 
         private final InvariantDeviceProfile mIdp;
-        private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
         private final ConcurrentLinkedQueue<LauncherIconsForPreview> mIconPool =
                 new ConcurrentLinkedQueue<>();
 
-        private boolean mDestroyed = false;
-
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base);
+            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE,
+                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
+                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE);
             mIdp = idp;
             mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
             mObjectMap.put(LauncherAppState.INSTANCE,
@@ -145,37 +136,6 @@
 
         }
 
-        @Override
-        public Context getApplicationContext() {
-            return this;
-        }
-
-        public void onDestroy() {
-            CustomWidgetManager.INSTANCE.get(this).onDestroy();
-            LauncherAppState.INSTANCE.get(this).onTerminate();
-            mDestroyed = true;
-        }
-
-        /**
-         * Find a cached object from mObjectMap if we have already created one. If not, generate
-         * an object using the provider.
-         */
-        public <T> T getObject(MainThreadInitializedObject<T> mainThreadInitializedObject,
-                MainThreadInitializedObject.ObjectProvider<T> provider) {
-            if (FeatureFlags.IS_STUDIO_BUILD && mDestroyed) {
-                throw new RuntimeException("Context already destroyed");
-            }
-            if (!mAllowedObjects.contains(mainThreadInitializedObject)) {
-                throw new IllegalStateException("Leaking unknown objects");
-            }
-            if (mObjectMap.containsKey(mainThreadInitializedObject)) {
-                return (T) mObjectMap.get(mainThreadInitializedObject);
-            }
-            T t = provider.get(this);
-            mObjectMap.put(mainThreadInitializedObject, t);
-            return t;
-        }
-
         public LauncherIcons newLauncherIcons(Context context, boolean shapeDetection) {
             LauncherIconsForPreview launcherIconsForPreview = mIconPool.poll();
             if (launcherIconsForPreview != null) {
diff --git a/src/com/android/launcher3/icons/IconCache.java b/src/com/android/launcher3/icons/IconCache.java
index cd13cd0..1a468ae 100644
--- a/src/com/android/launcher3/icons/IconCache.java
+++ b/src/com/android/launcher3/icons/IconCache.java
@@ -134,6 +134,9 @@
      * Closes the cache DB. This will clear any in-memory cache.
      */
     public void close() {
+        // This will clear all pending updates
+        getUpdateHandler();
+
         mIconDb.close();
     }
 
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index e2c0a32..94f29db 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -55,7 +55,7 @@
  * Utility class to cache properties of default display to avoid a system RPC on every call.
  */
 @SuppressLint("NewApi")
-public class DisplayController implements DisplayListener, ComponentCallbacks {
+public class DisplayController implements DisplayListener, ComponentCallbacks, SafeCloseable {
 
     private static final String TAG = "DisplayController";
 
@@ -79,6 +79,7 @@
 
     private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
     private Info mInfo;
+    private boolean mDestroyed = false;
 
     private DisplayController(Context context) {
         mContext = context;
@@ -111,6 +112,17 @@
     }
 
     @Override
+    public void close() {
+        mDestroyed = true;
+        if (mWindowContext != null) {
+            mWindowContext.unregisterComponentCallbacks(this);
+        } else {
+            // TODO: unregister broadcast receiver
+        }
+        mDM.unregisterDisplayListener(this);
+    }
+
+    @Override
     public final void onDisplayAdded(int displayId) { }
 
     @Override
@@ -157,6 +169,9 @@
      * Only used for pre-S
      */
     private void onConfigChanged(Intent intent) {
+        if (mDestroyed) {
+            return;
+        }
         Configuration config = mContext.getResources().getConfiguration();
         if (mInfo.fontScale != config.fontScale || mInfo.densityDpi != config.densityDpi) {
             Log.d(TAG, "Configuration changed, notifying listeners");
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index f6003dd..ef160b1 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -18,13 +18,20 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.os.Looper;
 
+import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
 
-import com.android.launcher3.graphics.LauncherPreviewRenderer.PreviewContext;
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -40,8 +47,8 @@
     }
 
     public T get(Context context) {
-        if (context instanceof PreviewContext) {
-            return ((PreviewContext) context).getObject(this, mProvider);
+        if (context instanceof SandboxContext) {
+            return ((SandboxContext) context).getObject(this, mProvider);
         }
 
         if (mValue == null) {
@@ -80,4 +87,78 @@
 
         T get(Context context);
     }
+
+    /**
+     * Abstract Context which allows custom implementations for
+     * {@link MainThreadInitializedObject} providers
+     */
+    public static abstract class SandboxContext extends ContextWrapper {
+
+        protected final Set<MainThreadInitializedObject> mAllowedObjects;
+        protected final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
+        protected final ArrayList<Object> mOrderedObjects = new ArrayList<>();
+
+        private final Object mDestroyLock = new Object();
+        private boolean mDestroyed = false;
+
+        public SandboxContext(Context base, MainThreadInitializedObject... allowedObjects) {
+            super(base);
+            mAllowedObjects = new HashSet<>(Arrays.asList(allowedObjects));
+        }
+
+        @Override
+        public Context getApplicationContext() {
+            return this;
+        }
+
+        public void onDestroy() {
+            synchronized (mDestroyLock) {
+                // Destroy in reverse order
+                for (int i = mOrderedObjects.size() - 1; i >= 0; i--) {
+                    Object o = mOrderedObjects.get(i);
+                    if (o instanceof SafeCloseable) {
+                        ((SafeCloseable) o).close();
+                    }
+                }
+                mDestroyed = true;
+            }
+        }
+
+        /**
+         * Find a cached object from mObjectMap if we have already created one. If not, generate
+         * an object using the provider.
+         */
+        private <T> T getObject(MainThreadInitializedObject<T> object, ObjectProvider<T> provider) {
+            synchronized (mDestroyLock) {
+                if (mDestroyed) {
+                    throw new RuntimeException("Context already destroyed");
+                }
+                if (!mAllowedObjects.contains(object)) {
+                    throw new IllegalStateException(
+                            "Leaking unknown objects " + object + "  " + provider);
+                }
+                T t = (T) mObjectMap.get(object);
+                if (t != null) {
+                    return t;
+                }
+                if (Looper.myLooper() == Looper.getMainLooper()) {
+                    t = createObject(provider);
+                    mObjectMap.put(object, t);
+                    mOrderedObjects.add(t);
+                    return t;
+                }
+            }
+
+            try {
+                return MAIN_EXECUTOR.submit(() -> getObject(object, provider)).get();
+            } catch (InterruptedException | ExecutionException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @UiThread
+        protected <T> T createObject(ObjectProvider<T> provider) {
+            return provider.get(this);
+        }
+    }
 }
diff --git a/src/com/android/launcher3/util/SettingsCache.java b/src/com/android/launcher3/util/SettingsCache.java
index 10611c7..0c5b722 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -47,7 +47,7 @@
  *
  * Cache will also be updated if a key queried is missing (even if it has no listeners registered).
  */
-public class SettingsCache extends ContentObserver {
+public class SettingsCache extends ContentObserver implements SafeCloseable {
 
     /** Hidden field Settings.Secure.NOTIFICATION_BADGING */
     public static final Uri NOTIFICATION_BADGING_URI =
@@ -69,7 +69,6 @@
     private final Map<Uri, CopyOnWriteArrayList<OnChangeListener>> mListenerMap = new HashMap<>();
     protected final ContentResolver mResolver;
 
-
     /**
      * Singleton instance
      */
@@ -82,6 +81,11 @@
     }
 
     @Override
+    public void close() {
+        mResolver.unregisterContentObserver(this);
+    }
+
+    @Override
     public void onChange(boolean selfChange, Uri uri) {
         // We use default of 1, but if we're getting an onChange call, can assume a non-default
         // value will exist
diff --git a/src/com/android/launcher3/util/UiThreadHelper.java b/src/com/android/launcher3/util/UiThreadHelper.java
index 0f40179..ac5368c 100644
--- a/src/com/android/launcher3/util/UiThreadHelper.java
+++ b/src/com/android/launcher3/util/UiThreadHelper.java
@@ -28,7 +28,7 @@
 import android.view.View;
 import android.view.inputmethod.InputMethodManager;
 
-import com.android.launcher3.Launcher;
+import com.android.launcher3.BaseActivity;
 import com.android.launcher3.views.ActivityContext;
 
 /**
@@ -56,7 +56,7 @@
                 STATS_LOGGER_KEY,
                 Message.obtain(
                         HANDLER.get(root.getContext()),
-                        () -> Launcher.cast(activityContext)
+                        () -> BaseActivity.fromContext(root.getContext())
                                 .getStatsLogManager()
                                 .logger()
                                 .log(LAUNCHER_ALLAPPS_KEYBOARD_CLOSED)
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 329a444..2e2a968 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -33,6 +33,7 @@
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PackageUserKey;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.systemui.plugins.CustomWidgetPlugin;
@@ -46,7 +47,7 @@
 /**
  * CustomWidgetManager handles custom widgets implemented as a plugin.
  */
-public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin> {
+public class CustomWidgetManager implements PluginListener<CustomWidgetPlugin>, SafeCloseable {
 
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
@@ -71,7 +72,8 @@
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
     }
 
-    public void onDestroy() {
+    @Override
+    public void close() {
         PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
     }