Revert "Revert "Migrating all model tests to Instrumentation tests""

This reverts commit 7a4a30d86d471e6c45adc2a9907efb27e9b1799b.

Test: Presubmit
Reason for revert: Fixing original bug

Bug: 196825541
Change-Id: Id4b1eb24a89564d264266d305aebea52917dfcd9

1 files changed, 87 insertions, 3 deletions

diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index f6003dd7bf..badcd35c23 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -18,13 +18,21 @@ package com.android.launcher3.util;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.content.Context;
+import android.content.ContextWrapper;
 import android.os.Looper;
+import android.util.Log;
 
+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 +48,8 @@ public class MainThreadInitializedObject<T> {
     }
 
     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 +88,80 @@ public class MainThreadInitializedObject<T> {
 
         T get(Context context);
     }
+
+    /**
+     * Abstract Context which allows custom implementations for
+     * {@link MainThreadInitializedObject} providers
+     */
+    public static abstract class SandboxContext extends ContextWrapper {
+
+        private static final String TAG = "SandboxContext";
+
+        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) {
+                    Log.e(TAG, "Static object access with a destroyed context");
+                }
+                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);
+        }
+    }
 }