Moving LauncherAppWidgetHolder to dagger

This would allow customizing the widget holder in LauncherPreview

Bug: 361850561
Test: Updated tests and presubmit
Flag: EXEMPT dagger
Change-Id: I32491169188992453693048986c57cb780fdf1d8
diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index 1ce28cf..49cee0f 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -23,7 +23,6 @@
 
     <string name="stats_log_manager_class" translatable="false">com.android.quickstep.logging.StatsLogCompatManager</string>
     <string name="test_information_handler_class" translatable="false">com.android.quickstep.QuickstepTestInformationHandler</string>
-    <string name="widget_holder_factory_class" translatable="false">com.android.launcher3.uioverrides.QuickstepWidgetHolder$QuickstepHolderFactory</string>
     <string name="instant_app_resolver_class" translatable="false">com.android.quickstep.InstantAppResolverImpl</string>
     <string name="app_launch_tracker_class" translatable="false">com.android.launcher3.appprediction.PredictionAppTracker</string>
     <string name="main_process_initializer_class" translatable="false">com.android.quickstep.QuickstepProcessInitializer</string>
diff --git a/quickstep/src/com/android/launcher3/dagger/Modules.kt b/quickstep/src/com/android/launcher3/dagger/Modules.kt
index 52be413..7671a82 100644
--- a/quickstep/src/com/android/launcher3/dagger/Modules.kt
+++ b/quickstep/src/com/android/launcher3/dagger/Modules.kt
@@ -16,11 +16,13 @@
 
 package com.android.launcher3.dagger
 
+import com.android.launcher3.uioverrides.QuickstepWidgetHolder.QuickstepWidgetHolderFactory
 import com.android.launcher3.uioverrides.SystemApiWrapper
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapperImpl
 import com.android.launcher3.util.ApiWrapper
 import com.android.launcher3.util.PluginManagerWrapper
 import com.android.launcher3.util.window.WindowManagerProxy
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactory
 import com.android.quickstep.util.GestureExclusionManager
 import com.android.quickstep.util.SystemWindowManagerProxy
 import dagger.Binds
@@ -40,6 +42,13 @@
 }
 
 @Module
+abstract class WidgetModule {
+
+    @Binds
+    abstract fun bindWidgetHolderFactory(factor: QuickstepWidgetHolderFactory): WidgetHolderFactory
+}
+
+@Module
 abstract class PluginManagerWrapperModule {
     @Binds
     abstract fun bindPluginManagerWrapper(impl: PluginManagerWrapperImpl): PluginManagerWrapper
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java
deleted file mode 100644
index 45813ce..0000000
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHost.java
+++ /dev/null
@@ -1,74 +0,0 @@
-/**
- * Copyright (C) 2022 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.uioverrides;
-
-import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
-
-import android.appwidget.AppWidgetHost;
-import android.appwidget.AppWidgetProviderInfo;
-import android.content.Context;
-import android.os.Looper;
-
-import androidx.annotation.NonNull;
-
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.util.Executors;
-import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
-import com.android.launcher3.widget.LauncherWidgetHolder;
-
-import java.util.function.IntConsumer;
-
-/**
- * {@link AppWidgetHost} that is used to receive the changes to the widgets without
- * storing any {@code Activity} info like that of the launcher.
- */
-final class QuickstepAppWidgetHost extends AppWidgetHost {
-    private final @NonNull Context mContext;
-    private final @NonNull IntConsumer mAppWidgetRemovedCallback;
-    private final @NonNull LauncherWidgetHolder.ProviderChangedListener mProvidersChangedListener;
-
-    QuickstepAppWidgetHost(@NonNull Context context, @NonNull IntConsumer appWidgetRemovedCallback,
-            @NonNull LauncherWidgetHolder.ProviderChangedListener listener,
-            @NonNull Looper looper) {
-        super(context, APPWIDGET_HOST_ID, null, looper);
-        mContext = context;
-        mAppWidgetRemovedCallback = appWidgetRemovedCallback;
-        mProvidersChangedListener = listener;
-    }
-
-    @Override
-    protected void onProvidersChanged() {
-        mProvidersChangedListener.notifyWidgetProvidersChanged();
-    }
-
-    @Override
-    public void onAppWidgetRemoved(int appWidgetId) {
-        // Route the call via model thread, in case it comes while a loader-bind is in progress
-        Executors.MODEL_EXECUTOR.execute(
-                () -> Executors.MAIN_EXECUTOR.execute(
-                        () -> mAppWidgetRemovedCallback.accept(appWidgetId)));
-    }
-
-    @Override
-    protected void onProviderChanged(int appWidgetId, @NonNull AppWidgetProviderInfo appWidget) {
-        LauncherAppWidgetProviderInfo info = LauncherAppWidgetProviderInfo.fromProviderInfo(
-                mContext, appWidget);
-        super.onProviderChanged(appWidgetId, info);
-        // The super method updates the dimensions of the providerInfo. Update the
-        // launcher spans accordingly.
-        info.initSpans(mContext, LauncherAppState.getIDP(mContext));
-    }
-}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHostProvider.kt b/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHostProvider.kt
new file mode 100644
index 0000000..1387cb7
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepAppWidgetHostProvider.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2025 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.uioverrides
+
+import android.app.ActivityThread
+import android.content.Context
+import android.content.ContextWrapper
+import com.android.launcher3.BuildConfig
+import com.android.launcher3.util.LooperExecutor
+import com.android.launcher3.widget.LauncherWidgetHolder
+import com.android.launcher3.widget.ListenableAppWidgetHost
+
+object QuickstepAppWidgetHostProvider {
+
+    /** Static widget host which is always listening and is lazily created */
+    @JvmStatic
+    val staticQuickstepHost: ListenableAppWidgetHost by lazy {
+        ListenableAppWidgetHost(
+                LooperContext(
+                    ActivityThread.currentApplication(),
+                    ListenableAppWidgetHost.widgetHolderExecutor,
+                ),
+                LauncherWidgetHolder.APPWIDGET_HOST_ID,
+            )
+            .apply { if (BuildConfig.WIDGETS_ENABLED) startListening() }
+    }
+
+    private class LooperContext(ctx: Context, val executor: LooperExecutor) : ContextWrapper(ctx) {
+
+        override fun getMainLooper() = executor.looper
+
+        override fun getMainExecutor() = executor
+    }
+}
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepInteractionHandler.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepInteractionHandler.java
index 26a1322..2f61eab 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepInteractionHandler.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepInteractionHandler.java
@@ -34,8 +34,11 @@
 import com.android.launcher3.util.ActivityOptionsWrapper;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 
+import java.util.function.Consumer;
+
 /** Provides a Quickstep specific animation when launching an activity from an app widget. */
-class QuickstepInteractionHandler implements RemoteViews.InteractionHandler {
+class QuickstepInteractionHandler implements RemoteViews.InteractionHandler,
+        Consumer<LauncherAppWidgetHostView> {
 
     private static final String TAG = "QuickstepInteractionHandler";
 
@@ -45,6 +48,11 @@
         mLauncher = launcher;
     }
 
+    @Override
+    public void accept(LauncherAppWidgetHostView host) {
+        host.setInteractionHandler(this);
+    }
+
     @SuppressWarnings("NewApi")
     @Override
     public boolean onInteraction(View view, PendingIntent pendingIntent,
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 806b8ab..605fd31 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -143,7 +143,6 @@
 import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
-import com.android.launcher3.uioverrides.QuickstepWidgetHolder.QuickstepHolderFactory;
 import com.android.launcher3.uioverrides.states.QuickstepAtomicAnimationFactory;
 import com.android.launcher3.uioverrides.touchcontrollers.NavBarToHomeTouchController;
 import com.android.launcher3.uioverrides.touchcontrollers.NoButtonNavbarToOverviewTouchController;
@@ -172,7 +171,6 @@
 import com.android.launcher3.util.StartActivityParams;
 import com.android.launcher3.util.TouchController;
 import com.android.launcher3.views.FloatingIconView;
-import com.android.launcher3.widget.LauncherWidgetHolder;
 import com.android.quickstep.OverviewCommandHelper;
 import com.android.quickstep.OverviewComponentObserver;
 import com.android.quickstep.OverviewComponentObserver.OverviewChangeListener;
@@ -298,6 +296,7 @@
 
     @Override
     protected void setupViews() {
+        getAppWidgetHolder().setOnViewCreationCallback(new QuickstepInteractionHandler(this));
         super.setupViews();
 
         mActionsView = findViewById(R.id.overview_actions_view);
@@ -726,15 +725,6 @@
     }
 
     @Override
-    protected LauncherWidgetHolder createAppWidgetHolder() {
-        final QuickstepHolderFactory factory =
-                (QuickstepHolderFactory) LauncherWidgetHolder.HolderFactory.newFactory(this);
-        return factory.newInstance(this,
-                appWidgetId -> getWorkspace().removeWidget(appWidgetId),
-                new QuickstepInteractionHandler(this));
-    }
-
-    @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
         if (savedInstanceState != null) {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
index 56fc4d1..9970a7d 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
@@ -16,13 +16,14 @@
 package com.android.launcher3.uioverrides;
 
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
+import static com.android.launcher3.uioverrides.QuickstepAppWidgetHostProvider.getStaticQuickstepHost;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.widget.ListenableAppWidgetHost.getWidgetHolderExecutor;
 
 import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
-import android.util.Log;
 import android.util.SparseArray;
 import android.widget.RemoteViews;
 
@@ -31,28 +32,26 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
-import java.util.ArrayList;
+import dagger.assisted.Assisted;
+import dagger.assisted.AssistedFactory;
+import dagger.assisted.AssistedInject;
+
 import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.function.BiConsumer;
-import java.util.function.IntConsumer;
 
 /**
  * {@link LauncherWidgetHolder} that puts the app widget host in the background
  */
 public final class QuickstepWidgetHolder extends LauncherWidgetHolder {
 
-    private static final String TAG = "QuickstepWidgetHolder";
-
     private static final UpdateKey<AppWidgetProviderInfo> KEY_PROVIDER_UPDATE =
             AppWidgetHostView::onUpdateProviderInfo;
     private static final UpdateKey<RemoteViews> KEY_VIEWS_UPDATE =
@@ -60,51 +59,17 @@
     private static final UpdateKey<Integer> KEY_VIEW_DATA_CHANGED =
             AppWidgetHostView::onViewDataChanged;
 
-    private static final List<QuickstepWidgetHolder> sHolders = new ArrayList<>();
     private static final SparseArray<QuickstepWidgetHolderListener> sListeners =
             new SparseArray<>();
 
-    private static AppWidgetHost sWidgetHost = null;
-
     private final UpdateHandler mUpdateHandler = this::onWidgetUpdate;
-    private final @Nullable RemoteViews.InteractionHandler mInteractionHandler;
-
-    private final @NonNull IntConsumer mAppWidgetRemovedCallback;
 
     // Map to all pending updated keyed with appWidgetId;
     private final SparseArray<PendingUpdate> mPendingUpdateMap = new SparseArray<>();
 
-    private QuickstepWidgetHolder(@NonNull Context context,
-            @Nullable IntConsumer appWidgetRemovedCallback,
-            @Nullable RemoteViews.InteractionHandler interactionHandler) {
-        super(context, appWidgetRemovedCallback);
-        mAppWidgetRemovedCallback = appWidgetRemovedCallback != null ? appWidgetRemovedCallback
-                : i -> {};
-        mInteractionHandler = interactionHandler;
-        MAIN_EXECUTOR.execute(() -> sHolders.add(this));
-    }
-
-    @Override
-    @NonNull
-    protected AppWidgetHost createHost(@NonNull Context context,
-            @Nullable IntConsumer appWidgetRemovedCallback) {
-        if (sWidgetHost == null) {
-            sWidgetHost = new QuickstepAppWidgetHost(context.getApplicationContext(),
-                    i -> MAIN_EXECUTOR.execute(() ->
-                            sHolders.forEach(h -> h.mAppWidgetRemovedCallback.accept(i))),
-                    () -> MAIN_EXECUTOR.execute(() ->
-                            sHolders.forEach(h ->
-                                    // Listeners might remove themselves from the list during the
-                                    // iteration. Creating a copy of the list to avoid exceptions
-                                    // for concurrent modification.
-                                    new ArrayList<>(h.mProviderChangedListeners).forEach(
-                                    ProviderChangedListener::notifyWidgetProvidersChanged))),
-                    getWidgetHolderExecutor().getLooper());
-            if (WIDGETS_ENABLED) {
-                sWidgetHost.startListening();
-            }
-        }
-        return sWidgetHost;
+    @AssistedInject
+    public QuickstepWidgetHolder(@Assisted("UI_CONTEXT") @NonNull Context context) {
+        super(context, getStaticQuickstepHost());
     }
 
     @Override
@@ -168,21 +133,6 @@
         sListeners.remove(appWidgetId);
     }
 
-    /**
-     * Called when the launcher is destroyed
-     */
-    @Override
-    public void destroy() {
-        try {
-            MAIN_EXECUTOR.submit(() -> {
-                clearViews();
-                sHolders.remove(this);
-            }).get();
-        } catch (Exception e) {
-            Log.e(TAG, "Failed to remove self from holder list", e);
-        }
-    }
-
     @Override
     protected boolean shouldListen(int flags) {
         return (flags & (FLAG_STATE_IS_NORMAL | FLAG_ACTIVITY_STARTED))
@@ -199,7 +149,7 @@
         }
 
         getWidgetHolderExecutor().execute(() -> {
-            sWidgetHost.setAppWidgetHidden();
+            mWidgetHost.setAppWidgetHidden();
             setListeningFlag(false);
         });
     }
@@ -239,7 +189,6 @@
     protected LauncherAppWidgetHostView createViewInternal(
             int appWidgetId, @NonNull LauncherAppWidgetProviderInfo appWidget) {
         LauncherAppWidgetHostView widgetView = new LauncherAppWidgetHostView(mContext);
-        widgetView.setInteractionHandler(mInteractionHandler);
         widgetView.setAppWidget(appWidgetId, appWidget);
         widgetView.updateAppWidget(getHolderListener(appWidgetId).addHolder(mUpdateHandler));
         return widgetView;
@@ -249,7 +198,7 @@
         QuickstepWidgetHolderListener listener = sListeners.get(appWidgetId);
         if (listener == null) {
             listener = new QuickstepWidgetHolderListener(appWidgetId);
-            sWidgetHost.setListener(appWidgetId, listener);
+            getStaticQuickstepHost().setListener(appWidgetId, listener);
             sListeners.put(appWidgetId, listener);
         }
         return listener;
@@ -322,44 +271,13 @@
         }
     }
 
-    /**
-     * {@code HolderFactory} subclass that takes an interaction handler as one of the parameters
-     * when creating a new instance.
-     */
-    public static class QuickstepHolderFactory extends HolderFactory {
 
-        @SuppressWarnings("unused")
-        public QuickstepHolderFactory(Context context) { }
+    /** A factory that generates new instances of {@code LauncherWidgetHolder} */
+    @AssistedFactory
+    public interface QuickstepWidgetHolderFactory extends WidgetHolderFactory {
 
         @Override
-        public LauncherWidgetHolder newInstance(@NonNull Context context,
-                @Nullable IntConsumer appWidgetRemovedCallback) {
-            return newInstance(context, appWidgetRemovedCallback, null);
-        }
-
-        /**
-         * @param context The context of the caller
-         * @param appWidgetRemovedCallback The callback that is called when widgets are removed
-         * @param interactionHandler The interaction handler when the widgets are clicked
-         * @return A new {@link LauncherWidgetHolder} instance
-         */
-        public LauncherWidgetHolder newInstance(@NonNull Context context,
-                @Nullable IntConsumer appWidgetRemovedCallback,
-                @Nullable RemoteViews.InteractionHandler interactionHandler) {
-
-            if (!FeatureFlags.ENABLE_WIDGET_HOST_IN_BACKGROUND.get()) {
-                return new LauncherWidgetHolder(context, appWidgetRemovedCallback) {
-                    @Override
-                    protected AppWidgetHost createHost(Context context,
-                            @Nullable IntConsumer appWidgetRemovedCallback) {
-                        AppWidgetHost host = super.createHost(context, appWidgetRemovedCallback);
-                        host.setInteractionHandler(interactionHandler);
-                        return host;
-                    }
-                };
-            }
-            return new QuickstepWidgetHolder(context, appWidgetRemovedCallback, interactionHandler);
-        }
+        QuickstepWidgetHolder newInstance(@Assisted("UI_CONTEXT") @NonNull Context context);
     }
 
     private interface UpdateKey<T> extends BiConsumer<AppWidgetHostView, T> { }
diff --git a/res/values/config.xml b/res/values/config.xml
index 3b48c9e..fc636a5 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -68,7 +68,6 @@
     <string name="app_launch_tracker_class" translatable="false"></string>
     <string name="test_information_handler_class" translatable="false"></string>
     <string name="secondary_display_predictions_class" translatable="false"></string>
-    <string name="widget_holder_factory_class" translatable="false"></string>
     <string name="taskbar_search_session_controller_class" translatable="false"></string>
     <string name="taskbar_model_callbacks_factory_class" translatable="false"></string>
     <string name="taskbar_view_callbacks_factory_class" translatable="false"></string>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 5c9392d..d5b3ed5 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -532,11 +532,14 @@
         mAllAppsController = new AllAppsTransitionController(this);
         mStateManager = new StateManager<>(this, NORMAL);
 
+        mAppWidgetManager = new WidgetManagerHelper(this);
+        mAppWidgetHolder = LauncherWidgetHolder.newInstance(this);
+        mAppWidgetHolder.setAppWidgetRemovedCallback(
+                appWidgetId -> getWorkspace().removeWidget(appWidgetId));
+
         setupViews();
         updateDisallowBack();
 
-        mAppWidgetManager = new WidgetManagerHelper(this);
-        mAppWidgetHolder = createAppWidgetHolder();
         mAppWidgetHolder.startListening();
         mAppWidgetHolder.addProviderChangeListener(() -> refreshAndBindWidgetsForPackageUser(null));
         mItemInflater = new ItemInflater<>(this, mAppWidgetHolder, getItemOnClickListener(),
@@ -1614,11 +1617,6 @@
         return instance;
     }
 
-    protected LauncherWidgetHolder createAppWidgetHolder() {
-        return LauncherWidgetHolder.HolderFactory.newFactory(this).newInstance(
-                this, appWidgetId -> getWorkspace().removeWidget(appWidgetId));
-    }
-
     @Override
     protected void onNewIntent(Intent intent) {
         if (Utilities.isRunningInTestHarness()) {
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 44dcc06..d987841 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -163,11 +163,6 @@
             "ENABLE_WIDGET_TRANSITION_FOR_RESIZING", DISABLED,
             "Enable widget transition animation when resizing the widgets");
 
-    // TODO(Block 25): Clean up flags
-    public static final BooleanFlag ENABLE_WIDGET_HOST_IN_BACKGROUND = getDebugFlag(270394384,
-            "ENABLE_WIDGET_HOST_IN_BACKGROUND", ENABLED,
-            "Enable background widget updates listening for widget holder");
-
     // TODO(Block 27): Clean up flags
     public static final BooleanFlag ENABLE_OVERLAY_CONNECTION_OPTIM = getDebugFlag(270392629,
             "ENABLE_OVERLAY_CONNECTION_OPTIM", DISABLED,
diff --git a/src/com/android/launcher3/dagger/LauncherAppModule.java b/src/com/android/launcher3/dagger/LauncherAppModule.java
index c58a414..0fd3219 100644
--- a/src/com/android/launcher3/dagger/LauncherAppModule.java
+++ b/src/com/android/launcher3/dagger/LauncherAppModule.java
@@ -23,6 +23,7 @@
         ApiWrapperModule.class,
         PluginManagerWrapperModule.class,
         StaticObjectModule.class,
+        WidgetModule.class,
         AppModule.class
 })
 public class LauncherAppModule {
diff --git a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
index c499097..f86772e 100644
--- a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
+++ b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
@@ -46,6 +46,7 @@
 import com.android.launcher3.util.WallpaperColorHints;
 import com.android.launcher3.util.window.RefreshRateTracker;
 import com.android.launcher3.util.window.WindowManagerProxy;
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactory;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import dagger.BindsInstance;
@@ -89,6 +90,7 @@
     WidgetsFilterDataProvider getWidgetsFilterDataProvider();
 
     LoaderCursorFactory getLoaderCursorFactory();
+    WidgetHolderFactory getWidgetHolderFactory();
 
     /** Builder for LauncherBaseAppComponent. */
     interface Builder {
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
index 91b899c..63d2954 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
@@ -16,8 +16,6 @@
 
 package com.android.launcher3.widget;
 
-import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
-
 import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.Context;
@@ -26,49 +24,18 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.util.Executors;
-import com.android.launcher3.widget.LauncherWidgetHolder.ProviderChangedListener;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.function.IntConsumer;
-
 /**
  * Specific {@link AppWidgetHost} that creates our {@link LauncherAppWidgetHostView}
  * which correctly captures all long-press events. This ensures that users can
  * always pick up and move widgets.
  */
-class LauncherAppWidgetHost extends AppWidgetHost {
-    @NonNull
-    private final List<ProviderChangedListener> mProviderChangeListeners;
-
-    @NonNull
-    private final Context mContext;
-
-    @Nullable
-    private final IntConsumer mAppWidgetRemovedCallback;
+class LauncherAppWidgetHost extends ListenableAppWidgetHost {
 
     @Nullable
     private ListenableHostView mViewToRecycle;
 
-    public LauncherAppWidgetHost(@NonNull Context context,
-            @Nullable IntConsumer appWidgetRemovedCallback,
-            List<ProviderChangedListener> providerChangeListeners) {
-        super(context, APPWIDGET_HOST_ID);
-        mContext = context;
-        mAppWidgetRemovedCallback = appWidgetRemovedCallback;
-        mProviderChangeListeners = providerChangeListeners;
-    }
-
-    @Override
-    protected void onProvidersChanged() {
-        if (!mProviderChangeListeners.isEmpty()) {
-            for (LauncherWidgetHolder.ProviderChangedListener callback :
-                    new ArrayList<>(mProviderChangeListeners)) {
-                callback.notifyWidgetProvidersChanged();
-            }
-        }
+    LauncherAppWidgetHost(@NonNull Context context, int appWidgetId) {
+        super(context, appWidgetId);
     }
 
     /**
@@ -94,35 +61,6 @@
     }
 
     /**
-     * Called when the AppWidget provider for a AppWidget has been upgraded to a new apk.
-     */
-    @Override
-    protected void onProviderChanged(int appWidgetId, @NonNull AppWidgetProviderInfo appWidget) {
-        LauncherAppWidgetProviderInfo info = LauncherAppWidgetProviderInfo.fromProviderInfo(
-                mContext, appWidget);
-        super.onProviderChanged(appWidgetId, info);
-        // The super method updates the dimensions of the providerInfo. Update the
-        // launcher spans accordingly.
-        info.initSpans(mContext, LauncherAppState.getIDP(mContext));
-    }
-
-    /**
-     * Called on an appWidget is removed for a widgetId
-     *
-     * @param appWidgetId TODO: make this override when SDK is updated
-     */
-    @Override
-    public void onAppWidgetRemoved(int appWidgetId) {
-        if (mAppWidgetRemovedCallback == null) {
-            return;
-        }
-        // Route the call via model thread, in case it comes while a loader-bind is in progress
-        Executors.MODEL_EXECUTOR.execute(
-                () -> Executors.MAIN_EXECUTOR.execute(
-                        () -> mAppWidgetRemovedCallback.accept(appWidgetId)));
-    }
-
-    /**
      * The same as super.clearViews(), except with the scope exposed
      */
     @Override
diff --git a/src/com/android/launcher3/widget/LauncherWidgetHolder.java b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
index 78197e2..642f35a 100644
--- a/src/com/android/launcher3/widget/LauncherWidgetHolder.java
+++ b/src/com/android/launcher3/widget/LauncherWidgetHolder.java
@@ -20,10 +20,9 @@
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
 import static com.android.launcher3.Flags.enableWorkspaceInflation;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.widget.LauncherAppWidgetProviderInfo.fromProviderInfo;
+import static com.android.launcher3.widget.ListenableAppWidgetHost.getWidgetHolderExecutor;
 
-import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetHostView;
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
@@ -32,6 +31,7 @@
 import android.content.Intent;
 import android.os.Bundle;
 import android.os.Looper;
+import android.util.Log;
 import android.util.SparseArray;
 import android.widget.Toast;
 
@@ -43,18 +43,23 @@
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.dagger.LauncherComponentProvider;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
-import com.android.launcher3.util.LooperExecutor;
-import com.android.launcher3.util.ResourceBasedOverride;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.widget.ListenableAppWidgetHost.ProviderChangedListener;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
+import dagger.assisted.Assisted;
+import dagger.assisted.AssistedFactory;
+import dagger.assisted.AssistedInject;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 import java.util.function.IntConsumer;
 
 /**
@@ -62,51 +67,57 @@
  * background.
  */
 public class LauncherWidgetHolder {
+
+    private static final String TAG = "LauncherWidgetHolder";
+
     public static final int APPWIDGET_HOST_ID = 1024;
 
     protected static final int FLAG_LISTENING = 1;
     protected static final int FLAG_STATE_IS_NORMAL = 1 << 1;
     protected static final int FLAG_ACTIVITY_STARTED = 1 << 2;
     protected static final int FLAG_ACTIVITY_RESUMED = 1 << 3;
+
     private static final int FLAGS_SHOULD_LISTEN =
             FLAG_STATE_IS_NORMAL | FLAG_ACTIVITY_STARTED | FLAG_ACTIVITY_RESUMED;
 
-    @NonNull
-    protected final Context mContext;
-
-    @NonNull
-    private final AppWidgetHost mWidgetHost;
-
-    @NonNull
-    protected final SparseArray<LauncherAppWidgetHostView> mViews = new SparseArray<>();
-    protected final List<ProviderChangedListener> mProviderChangedListeners = new ArrayList<>();
-
-    protected AtomicInteger mFlags = new AtomicInteger(FLAG_STATE_IS_NORMAL);
-
     // TODO(b/191735836): Replace with ActivityOptions.KEY_SPLASH_SCREEN_STYLE when un-hidden
     private static final String KEY_SPLASH_SCREEN_STYLE = "android.activity.splashScreenStyle";
     // TODO(b/191735836): Replace with SplashScreen.SPLASH_SCREEN_STYLE_EMPTY when un-hidden
     private static final int SPLASH_SCREEN_STYLE_EMPTY = 0;
 
-    protected LauncherWidgetHolder(@NonNull Context context,
-            @Nullable IntConsumer appWidgetRemovedCallback) {
+    @NonNull
+    protected final Context mContext;
+
+    @NonNull
+    protected final ListenableAppWidgetHost mWidgetHost;
+
+    @NonNull
+    protected final SparseArray<LauncherAppWidgetHostView> mViews = new SparseArray<>();
+
+    /** package visibility */
+    final List<ProviderChangedListener> mProviderChangedListeners = new ArrayList<>();
+
+    protected AtomicInteger mFlags = new AtomicInteger(FLAG_STATE_IS_NORMAL);
+
+    @Nullable
+    private Consumer<LauncherAppWidgetHostView> mOnViewCreationCallback;
+
+    /** package visibility */
+    @Nullable IntConsumer mAppWidgetRemovedCallback;
+
+    @AssistedInject
+    protected LauncherWidgetHolder(@Assisted("UI_CONTEXT") @NonNull Context context) {
+        this(context, new LauncherAppWidgetHost(context, APPWIDGET_HOST_ID));
+    }
+
+    protected LauncherWidgetHolder(
+            @NonNull Context context, @NonNull ListenableAppWidgetHost appWidgetHost) {
         mContext = context;
-        mWidgetHost = createHost(context, appWidgetRemovedCallback);
+        mWidgetHost = appWidgetHost;
+        MAIN_EXECUTOR.execute(() ->  mWidgetHost.getHolders().add(this));
     }
 
-    protected AppWidgetHost createHost(
-            Context context, @Nullable IntConsumer appWidgetRemovedCallback) {
-        return new LauncherAppWidgetHost(
-                context, appWidgetRemovedCallback, mProviderChangedListeners);
-    }
-
-    protected LooperExecutor getWidgetHolderExecutor() {
-        return UI_HELPER_EXECUTOR;
-    }
-
-    /**
-     * Starts listening to the widget updates from the server side
-     */
+    /** Starts listening to the widget updates from the server side */
     public void startListening() {
         if (!WIDGETS_ENABLED) {
             return;
@@ -127,13 +138,11 @@
             // TODO: Investigate why widgetHost.startListening() always return non-empty updates
             setListeningFlag(true);
 
-            MAIN_EXECUTOR.execute(() -> updateDeferredView());
+            MAIN_EXECUTOR.execute(this::updateDeferredView);
         });
     }
 
-    /**
-     * Update any views which have been deferred because the host was not listening.
-     */
+    /** Update any views which have been deferred because the host was not listening */
     protected void updateDeferredView() {
         // Update any views which have been deferred because the host was not listening.
         // We go in reverse order and inflate any deferred or cached widget
@@ -180,7 +189,14 @@
      * Called when the launcher is destroyed
      */
     public void destroy() {
-        // No-op
+        try {
+            MAIN_EXECUTOR.submit(() -> {
+                clearViews();
+                mWidgetHost.getHolders().remove(this);
+            }).get();
+        } catch (Exception e) {
+            Log.e(TAG, "Failed to remove self from holder list", e);
+        }
     }
 
     /**
@@ -198,8 +214,7 @@
      * Add a listener that is triggered when the providers of the widgets are changed
      * @param listener The listener that notifies when the providers changed
      */
-    public void addProviderChangeListener(
-            @NonNull LauncherWidgetHolder.ProviderChangedListener listener) {
+    public void addProviderChangeListener(@NonNull ProviderChangedListener listener) {
         MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.add(listener));
     }
 
@@ -207,12 +222,23 @@
      * Remove the specified listener from the host
      * @param listener The listener that is to be removed from the host
      */
-    public void removeProviderChangeListener(
-            LauncherWidgetHolder.ProviderChangedListener listener) {
+    public void removeProviderChangeListener(ProviderChangedListener listener) {
         MAIN_EXECUTOR.execute(() -> mProviderChangedListeners.remove(listener));
     }
 
     /**
+     * Sets a callbacks for whenever a widget view is created
+     */
+    public void setOnViewCreationCallback(@Nullable Consumer<LauncherAppWidgetHostView> callback) {
+        mOnViewCreationCallback = callback;
+    }
+
+    /** Sets a callback for listening app widget removals */
+    public void setAppWidgetRemovedCallback(@Nullable IntConsumer callback) {
+        mAppWidgetRemovedCallback = callback;
+    }
+
+    /**
      * Starts the configuration activity for the widget
      * @param activity The activity in which to start the configuration page
      * @param widgetId The ID of the widget
@@ -284,9 +310,7 @@
         activity.startActivityForResult(intent, requestCode);
     }
 
-    /**
-     * Stop the host from listening to the widget updates
-     */
+    /** Stop the host from listening to the widget updates */
     public void stopListening() {
         if (!WIDGETS_ENABLED) {
             return;
@@ -298,8 +322,8 @@
     }
 
     /**
-     * Update {@link FLAG_LISTENING} on {@link mFlags} after making binder calls from
-     * {@link sWidgetHost}.
+     * Update {@link #FLAG_LISTENING} on {@link #mFlags} after making binder calls from
+     * {@link #mWidgetHost}.
      */
     @WorkerThread
     protected void setListeningFlag(final boolean isListening) {
@@ -350,6 +374,7 @@
         }
 
         LauncherAppWidgetHostView view = createViewInternal(appWidgetId, appWidget);
+        if (mOnViewCreationCallback != null) mOnViewCreationCallback.accept(view);
         // Do not update mViews on a background thread call, as the holder is not thread safe.
         if (!enableWorkspaceInflation() || Looper.myLooper() == Looper.getMainLooper()) {
             mViews.put(appWidgetId, view);
@@ -368,8 +393,8 @@
 
         // Binder can also inflate placeholder widgets in case of backup-restore. Skip
         // attaching such widgets
-        boolean isRealWidget = ((view instanceof PendingAppWidgetHostView pw)
-                ? pw.isDeferredWidget() : true)
+        boolean isRealWidget = (!(view instanceof PendingAppWidgetHostView pw)
+                || pw.isDeferredWidget())
                 && view.getAppWidgetInfo() != null;
         if (isRealWidget && mViews.get(view.getAppWidgetId()) != view) {
             view = recycleExistingView(view);
@@ -446,28 +471,13 @@
         }
     }
 
-    /**
-     * Listener for getting notifications on provider changes.
-     */
-    public interface ProviderChangedListener {
-        /**
-         * Notify the listener that the providers have changed
-         */
-        void notifyWidgetProvidersChanged();
-    }
-
-    /**
-     * Clears all the views from the host
-     */
+    /** Clears all the views from the host */
     public void clearViews() {
-        LauncherAppWidgetHost tempHost = (LauncherAppWidgetHost) mWidgetHost;
-        tempHost.clearViews();
+        ((LauncherAppWidgetHost) mWidgetHost).clearViews();
         mViews.clear();
     }
 
-    /**
-     * Clears all the internal widget views
-     */
+    /** Clears all the internal widget views */
     public void clearWidgetViews() {
         clearViews();
     }
@@ -514,32 +524,19 @@
      * Returns the new LauncherWidgetHolder instance
      */
     public static LauncherWidgetHolder newInstance(Context context) {
-        return HolderFactory.newFactory(context).newInstance(context, null);
+        return LauncherComponentProvider.get(context).getWidgetHolderFactory().newInstance(context);
     }
 
-    /**
-     * A factory class that generates new instances of {@code LauncherWidgetHolder}
-     */
-    public static class HolderFactory implements ResourceBasedOverride {
+    /** A factory that generates new instances of {@code LauncherWidgetHolder} */
+    public interface WidgetHolderFactory {
 
-        /**
-         * @param context The context of the caller
-         * @param appWidgetRemovedCallback The callback that is called when widgets are removed
-         * @return A new instance of {@code LauncherWidgetHolder}
-         */
-        public LauncherWidgetHolder newInstance(@NonNull Context context,
-                @Nullable IntConsumer appWidgetRemovedCallback) {
-            return new LauncherWidgetHolder(context, appWidgetRemovedCallback);
-        }
+        LauncherWidgetHolder newInstance(@NonNull Context context);
+    }
 
-        /**
-         * @param context The context of the caller
-         * @return A new instance of factory class for widget holders. If not specified, returning
-         * {@code HolderFactory} by default.
-         */
-        public static HolderFactory newFactory(Context context) {
-            return Overrides.getObject(
-                    HolderFactory.class, context, R.string.widget_holder_factory_class);
-        }
+    /** A factory that generates new instances of {@code LauncherWidgetHolder} */
+    @AssistedFactory
+    public interface WidgetHolderFactoryImpl extends WidgetHolderFactory {
+
+        LauncherWidgetHolder newInstance(@Assisted("UI_CONTEXT") @NonNull Context context);
     }
 }
diff --git a/src/com/android/launcher3/widget/ListenableAppWidgetHost.kt b/src/com/android/launcher3/widget/ListenableAppWidgetHost.kt
new file mode 100644
index 0000000..58bf0aa
--- /dev/null
+++ b/src/com/android/launcher3/widget/ListenableAppWidgetHost.kt
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2025 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.widget
+
+import android.appwidget.AppWidgetHost
+import android.appwidget.AppWidgetProviderInfo
+import android.content.Context
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.util.Executors
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+import com.android.launcher3.util.LooperExecutor
+
+open class ListenableAppWidgetHost(private val ctx: Context, hostId: Int) :
+    AppWidgetHost(ctx, hostId) {
+
+    protected val holders = mutableListOf<LauncherWidgetHolder>()
+
+    override fun onProvidersChanged() {
+        MAIN_EXECUTOR.execute {
+            holders.forEach { holder ->
+                // Listeners might remove themselves from the list during the iteration.
+                // Creating a copy of the list to avoid exceptions for concurrent modification.
+                holder.mProviderChangedListeners.toList().forEach {
+                    it.notifyWidgetProvidersChanged()
+                }
+            }
+        }
+    }
+
+    override fun onAppWidgetRemoved(appWidgetId: Int) {
+        // Route the call via model thread, in case it comes while a loader-bind is in progress
+        MODEL_EXECUTOR.execute {
+            MAIN_EXECUTOR.execute {
+                holders.forEach { it.mAppWidgetRemovedCallback?.accept(appWidgetId) }
+            }
+        }
+    }
+
+    override fun onProviderChanged(appWidgetId: Int, appWidget: AppWidgetProviderInfo) {
+        val info = LauncherAppWidgetProviderInfo.fromProviderInfo(ctx, appWidget)
+        super.onProviderChanged(appWidgetId, info)
+        // The super method updates the dimensions of the providerInfo. Update the
+        // launcher spans accordingly.
+        info.initSpans(ctx, InvariantDeviceProfile.INSTANCE.get(ctx))
+    }
+
+    /** Listener for getting notifications on provider changes. */
+    fun interface ProviderChangedListener {
+        /** Notify the listener that the providers have changed */
+        fun notifyWidgetProvidersChanged()
+    }
+
+    companion object {
+
+        @JvmStatic val widgetHolderExecutor: LooperExecutor = Executors.UI_HELPER_EXECUTOR
+    }
+}
diff --git a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
index cd8e457..1c29f89 100644
--- a/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/PendingAppWidgetHostView.java
@@ -66,7 +66,7 @@
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.Themes;
-import com.android.launcher3.widget.LauncherWidgetHolder.ProviderChangedListener;
+import com.android.launcher3.widget.ListenableAppWidgetHost.ProviderChangedListener;
 
 import java.util.List;
 
diff --git a/src_no_quickstep/com/android/launcher3/dagger/Modules.kt b/src_no_quickstep/com/android/launcher3/dagger/Modules.kt
index c3bf7c5..7dbe9c3 100644
--- a/src_no_quickstep/com/android/launcher3/dagger/Modules.kt
+++ b/src_no_quickstep/com/android/launcher3/dagger/Modules.kt
@@ -16,6 +16,9 @@
 
 package com.android.launcher3.dagger
 
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactory
+import com.android.launcher3.widget.LauncherWidgetHolder.WidgetHolderFactoryImpl
+import dagger.Binds
 import dagger.Module
 
 private object Modules {}
@@ -24,6 +27,12 @@
 
 @Module abstract class ApiWrapperModule {}
 
+@Module
+abstract class WidgetModule {
+    @Binds
+    abstract fun bindWidgetHolderFactory(factor: WidgetHolderFactoryImpl): WidgetHolderFactory
+}
+
 @Module abstract class PluginManagerWrapperModule {}
 
 @Module object StaticObjectModule {}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/DaggerGraphs.kt b/tests/multivalentTests/src/com/android/launcher3/util/DaggerGraphs.kt
index b66a9d3..a76ccf0 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/DaggerGraphs.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/DaggerGraphs.kt
@@ -21,6 +21,7 @@
 import com.android.launcher3.dagger.ApiWrapperModule
 import com.android.launcher3.dagger.AppModule
 import com.android.launcher3.dagger.StaticObjectModule
+import com.android.launcher3.dagger.WidgetModule
 import com.android.launcher3.dagger.WindowManagerProxyModule
 import dagger.Binds
 import dagger.Module
@@ -39,15 +40,27 @@
             ApiWrapperModule::class,
             WindowManagerProxyModule::class,
             StaticObjectModule::class,
+            WidgetModule::class,
             AppModule::class,
         ]
 )
 class AllModulesForTest
 
 /** All modules except the WMProxy */
-@Module(includes = [ApiWrapperModule::class, StaticObjectModule::class, AppModule::class])
+@Module(
+    includes =
+        [ApiWrapperModule::class, StaticObjectModule::class, AppModule::class, WidgetModule::class]
+)
 class AllModulesMinusWMProxy
 
 /** All modules except the ApiWrapper */
-@Module(includes = [WindowManagerProxyModule::class, StaticObjectModule::class, AppModule::class])
+@Module(
+    includes =
+        [
+            WindowManagerProxyModule::class,
+            StaticObjectModule::class,
+            AppModule::class,
+            WidgetModule::class,
+        ]
+)
 class AllModulesMinusApiWrapper
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt
index 79b493a..d5d1f4a 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt
@@ -21,32 +21,22 @@
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.util.ActivityContextWrapper
 import com.android.launcher3.util.Executors
+import com.android.launcher3.util.TestUtil
 import java.util.function.IntConsumer
 import org.junit.Assert.assertNotSame
 import org.junit.Assert.assertNull
 import org.junit.Assert.assertSame
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.verify
-import org.mockito.MockitoAnnotations
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class LauncherAppWidgetHostTest {
 
-    @Mock private lateinit var onAppWidgetRemovedCallback: IntConsumer
-
     private val context = ActivityContextWrapper(getInstrumentation().targetContext)
-    private lateinit var underTest: LauncherAppWidgetHost
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-        underTest = LauncherAppWidgetHost(context, onAppWidgetRemovedCallback, emptyList())
-    }
+    private var underTest = LauncherAppWidgetHost(context, HOST_ID)
 
     @Test
     fun `Host set view to recycle`() {
@@ -74,15 +64,20 @@
 
     @Test
     fun `Runnable called when app widget removed`() {
+        val holder = LauncherWidgetHolder(context, underTest)
+        holder.setAppWidgetRemovedCallback(mock(IntConsumer::class.java))
+        TestUtil.runOnExecutorSync(Executors.MAIN_EXECUTOR) {}
+
         underTest.onAppWidgetRemoved(WIDGET_ID)
 
         Executors.MODEL_EXECUTOR.submit {}.get()
         getInstrumentation().waitForIdleSync()
 
-        verify(onAppWidgetRemovedCallback).accept(WIDGET_ID)
+        verify(holder.mAppWidgetRemovedCallback!!).accept(WIDGET_ID)
     }
 
     companion object {
+        const val HOST_ID = 2233
         const val WIDGET_ID = 10001
     }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherWidgetHolderTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherWidgetHolderTest.kt
index 1a659e2..44f29d6 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherWidgetHolderTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherWidgetHolderTest.kt
@@ -47,7 +47,7 @@
     fun setUp() {
         assertTrue(WIDGETS_ENABLED)
         widgetHolder =
-            LauncherWidgetHolder(ActivityContextWrapper(getInstrumentation().targetContext)) {}
+            LauncherWidgetHolder(ActivityContextWrapper(getInstrumentation().targetContext))
     }
 
     @After
@@ -62,7 +62,7 @@
         widgetHolder.setListeningFlag(false)
         assertFalse(widgetHolder.isListening)
         widgetHolder.startListening()
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         getInstrumentation().waitForIdleSync()
         assertTrue(widgetHolder.isListening)
         verify(testView, times(1)).reInflate()
@@ -73,10 +73,10 @@
     fun holder_start_listening_after_activity_start() {
         widgetHolder.setShouldListenFlag(FLAG_STATE_IS_NORMAL or FLAG_ACTIVITY_RESUMED, true)
         widgetHolder.setActivityStarted(false)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertFalse(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
         widgetHolder.setActivityStarted(true)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertTrue(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
     }
 
@@ -84,10 +84,10 @@
     fun holder_start_listening_after_activity_resume() {
         widgetHolder.setShouldListenFlag(FLAG_STATE_IS_NORMAL or FLAG_ACTIVITY_STARTED, true)
         widgetHolder.setActivityResumed(false)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertFalse(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
         widgetHolder.setActivityResumed(true)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertTrue(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
     }
 
@@ -95,10 +95,10 @@
     fun holder_start_listening_after_state_normal() {
         widgetHolder.setShouldListenFlag(FLAG_ACTIVITY_RESUMED or FLAG_ACTIVITY_STARTED, true)
         widgetHolder.setStateIsNormal(false)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertFalse(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
         widgetHolder.setStateIsNormal(true)
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertTrue(widgetHolder.shouldListen(widgetHolder.mFlags.get()))
     }
 
@@ -117,7 +117,7 @@
 
     @Test
     fun holder_add_provider_change_listener() {
-        val listener = LauncherWidgetHolder.ProviderChangedListener {}
+        val listener = ListenableAppWidgetHost.ProviderChangedListener {}
         widgetHolder.addProviderChangeListener(listener)
         getInstrumentation().waitForIdleSync()
         assertEquals(1, widgetHolder.mProviderChangedListeners.size)
@@ -127,7 +127,7 @@
 
     @Test
     fun holder_remove_provider_change_listener() {
-        val listener = LauncherWidgetHolder.ProviderChangedListener {}
+        val listener = ListenableAppWidgetHost.ProviderChangedListener {}
         widgetHolder.addProviderChangeListener(listener)
         widgetHolder.removeProviderChangeListener(listener)
         getInstrumentation().waitForIdleSync()
@@ -139,7 +139,7 @@
         widgetHolder.setListeningFlag(true)
         assertTrue(widgetHolder.isListening)
         widgetHolder.stopListening()
-        widgetHolder.widgetHolderExecutor.submit {}.get()
+        ListenableAppWidgetHost.widgetHolderExecutor.submit {}.get()
         assertFalse(widgetHolder.isListening)
     }