diff --git a/src/com/android/launcher3/qsb/QsbContainerView.java b/src/com/android/launcher3/qsb/QsbContainerView.java
index b26d39f..8e29df1 100644
--- a/src/com/android/launcher3/qsb/QsbContainerView.java
+++ b/src/com/android/launcher3/qsb/QsbContainerView.java
@@ -16,6 +16,10 @@
 
 package com.android.launcher3.qsb;
 
+import static android.appwidget.AppWidgetManager.ACTION_APPWIDGET_BIND;
+import static android.appwidget.AppWidgetManager.EXTRA_APPWIDGET_ID;
+import static android.appwidget.AppWidgetManager.EXTRA_APPWIDGET_PROVIDER;
+
 import android.app.Activity;
 import android.app.Fragment;
 import android.app.SearchManager;
@@ -74,11 +78,12 @@
     /**
      * A fragment to display the QSB.
      */
-    public static class QsbFragment extends Fragment implements View.OnClickListener {
+    public static class QsbFragment extends Fragment {
 
+        public static final int QSB_WIDGET_HOST_ID = 1026;
         private static final int REQUEST_BIND_QSB = 1;
-        private static final String QSB_WIDGET_ID = "qsb_widget_id";
 
+        protected String mKeyWidgetId = "qsb_widget_id";
         private QsbWidgetHost mQsbWidgetHost;
         private AppWidgetProviderInfo mWidgetInfo;
         private QsbWidgetHostView mQsb;
@@ -90,10 +95,15 @@
         @Override
         public void onCreate(Bundle savedInstanceState) {
             super.onCreate(savedInstanceState);
-            mQsbWidgetHost = new QsbWidgetHost(getActivity());
+            mQsbWidgetHost = createHost();
             mOrientation = getContext().getResources().getConfiguration().orientation;
         }
 
+        protected QsbWidgetHost createHost() {
+            return new QsbWidgetHost(getActivity(), QSB_WIDGET_HOST_ID,
+                    (c) -> new QsbWidgetHostView(c));
+        }
+
         private FrameLayout mWrapper;
 
         @Override
@@ -110,24 +120,16 @@
         }
 
         private View createQsb(ViewGroup container) {
-            Activity activity = getActivity();
-            mWidgetInfo = getSearchWidgetProvider(activity);
+            mWidgetInfo = getSearchWidgetProvider();
             if (mWidgetInfo == null) {
                 // There is no search provider, just show the default widget.
-                return QsbWidgetHostView.getDefaultView(container);
+                return getDefaultView(container, false /* show setup icon */);
             }
-
+            Bundle opts = createBindOptions();
+            Activity activity = getActivity();
             AppWidgetManager widgetManager = AppWidgetManager.getInstance(activity);
-            InvariantDeviceProfile idp = LauncherAppState.getIDP(activity);
 
-            Bundle opts = new Bundle();
-            Rect size = AppWidgetResizeFrame.getWidgetSizeRanges(activity, idp.numColumns, 1, null);
-            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_WIDTH, size.left);
-            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_HEIGHT, size.top);
-            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_WIDTH, size.right);
-            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_HEIGHT, size.bottom);
-
-            int widgetId = Utilities.getPrefs(activity).getInt(QSB_WIDGET_ID, -1);
+            int widgetId = Utilities.getPrefs(activity).getInt(mKeyWidgetId, -1);
             AppWidgetProviderInfo widgetInfo = widgetManager.getAppWidgetInfo(widgetId);
             boolean isWidgetBound = (widgetInfo != null) &&
                     widgetInfo.provider.equals(mWidgetInfo.provider);
@@ -166,32 +168,18 @@
             }
 
             // Return a default widget with setup icon.
-            View v = QsbWidgetHostView.getDefaultView(container);
-            View setupButton = v.findViewById(R.id.btn_qsb_setup);
-            setupButton.setVisibility(View.VISIBLE);
-            setupButton.setOnClickListener(this);
-            return v;
+            return getDefaultView(container, true /* show setup icon */);
         }
 
         private void saveWidgetId(int widgetId) {
-            Utilities.getPrefs(getActivity()).edit().putInt(QSB_WIDGET_ID, widgetId).apply();
-        }
-
-        @Override
-        public void onClick(View view) {
-            // Start intent for bind the widget
-            Intent intent = new Intent(AppWidgetManager.ACTION_APPWIDGET_BIND);
-            // Allocate a new widget id for QSB
-            intent.putExtra(AppWidgetManager.EXTRA_APPWIDGET_ID, mQsbWidgetHost.allocateAppWidgetId());
-            intent.putExtra(AppWidgetManager.EXTRA_APPWIDGET_PROVIDER, mWidgetInfo.provider);
-            startActivityForResult(intent, REQUEST_BIND_QSB);
+            Utilities.getPrefs(getActivity()).edit().putInt(mKeyWidgetId, widgetId).apply();
         }
 
         @Override
         public void onActivityResult(int requestCode, int resultCode, Intent data) {
             if (requestCode == REQUEST_BIND_QSB) {
                 if (resultCode == Activity.RESULT_OK) {
-                    saveWidgetId(data.getIntExtra(AppWidgetManager.EXTRA_APPWIDGET_ID, -1));
+                    saveWidgetId(data.getIntExtra(EXTRA_APPWIDGET_ID, -1));
                     rebindFragment();
                 } else {
                     mQsbWidgetHost.deleteHost();
@@ -228,48 +216,83 @@
         public boolean isQsbEnabled() {
             return FeatureFlags.QSB_ON_FIRST_SCREEN;
         }
-    }
 
-    /**
-     * Returns a widget with category {@link AppWidgetProviderInfo#WIDGET_CATEGORY_SEARCHBOX}
-     * provided by the same package which is set to be global search activity.
-     * If widgetCategory is not supported, or no such widget is found, returns the first widget
-     * provided by the package.
-     */
-    public static AppWidgetProviderInfo getSearchWidgetProvider(Context context) {
-        SearchManager searchManager =
-                (SearchManager) context.getSystemService(Context.SEARCH_SERVICE);
-        ComponentName searchComponent = searchManager.getGlobalSearchActivity();
-        if (searchComponent == null) return null;
-        String providerPkg = searchComponent.getPackageName();
+        protected Bundle createBindOptions() {
+            InvariantDeviceProfile idp = LauncherAppState.getIDP(getActivity());
 
-        AppWidgetProviderInfo defaultWidgetForSearchPackage = null;
+            Bundle opts = new Bundle();
+            Rect size = AppWidgetResizeFrame.getWidgetSizeRanges(getActivity(),
+                    idp.numColumns, 1, null);
+            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_WIDTH, size.left);
+            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MIN_HEIGHT, size.top);
+            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_WIDTH, size.right);
+            opts.putInt(AppWidgetManager.OPTION_APPWIDGET_MAX_HEIGHT, size.bottom);
+            return opts;
+        }
 
-        AppWidgetManager appWidgetManager = AppWidgetManager.getInstance(context);
-        for (AppWidgetProviderInfo info : appWidgetManager.getInstalledProviders()) {
-            if (info.provider.getPackageName().equals(providerPkg) && info.configure == null) {
-                if ((info.widgetCategory & AppWidgetProviderInfo.WIDGET_CATEGORY_SEARCHBOX) != 0) {
-                    return info;
-                } else if (defaultWidgetForSearchPackage == null) {
-                    defaultWidgetForSearchPackage = info;
+        protected View getDefaultView(ViewGroup container, boolean showSetupIcon) {
+            // Return a default widget with setup icon.
+            View v = QsbWidgetHostView.getDefaultView(container);
+            if (showSetupIcon) {
+                View setupButton = v.findViewById(R.id.btn_qsb_setup);
+                setupButton.setVisibility(View.VISIBLE);
+                setupButton.setOnClickListener((v2) -> startActivityForResult(
+                        new Intent(ACTION_APPWIDGET_BIND)
+                                .putExtra(EXTRA_APPWIDGET_ID, mQsbWidgetHost.allocateAppWidgetId())
+                                .putExtra(EXTRA_APPWIDGET_PROVIDER, mWidgetInfo.provider),
+                        REQUEST_BIND_QSB));
+            }
+            return v;
+        }
+
+        /**
+         * Returns a widget with category {@link AppWidgetProviderInfo#WIDGET_CATEGORY_SEARCHBOX}
+         * provided by the same package which is set to be global search activity.
+         * If widgetCategory is not supported, or no such widget is found, returns the first widget
+         * provided by the package.
+         */
+        protected AppWidgetProviderInfo getSearchWidgetProvider() {
+            SearchManager searchManager =
+                    (SearchManager) getActivity().getSystemService(Context.SEARCH_SERVICE);
+            ComponentName searchComponent = searchManager.getGlobalSearchActivity();
+            if (searchComponent == null) return null;
+            String providerPkg = searchComponent.getPackageName();
+
+            AppWidgetProviderInfo defaultWidgetForSearchPackage = null;
+
+            AppWidgetManager appWidgetManager = AppWidgetManager.getInstance(getActivity());
+            for (AppWidgetProviderInfo info : appWidgetManager.getInstalledProviders()) {
+                if (info.provider.getPackageName().equals(providerPkg) && info.configure == null) {
+                    if ((info.widgetCategory
+                            & AppWidgetProviderInfo.WIDGET_CATEGORY_SEARCHBOX) != 0) {
+                        return info;
+                    } else if (defaultWidgetForSearchPackage == null) {
+                        defaultWidgetForSearchPackage = info;
+                    }
                 }
             }
+            return defaultWidgetForSearchPackage;
         }
-        return defaultWidgetForSearchPackage;
     }
 
-    private static class QsbWidgetHost extends AppWidgetHost {
+    public static class QsbWidgetHost extends AppWidgetHost {
 
-        private static final int QSB_WIDGET_HOST_ID = 1026;
+        private final WidgetViewFactory mViewFactory;
 
-        public QsbWidgetHost(Context context) {
-            super(context, QSB_WIDGET_HOST_ID);
+        public QsbWidgetHost(Context context, int hostId, WidgetViewFactory viewFactory) {
+            super(context, hostId);
+            mViewFactory = viewFactory;
         }
 
         @Override
         protected AppWidgetHostView onCreateView(
                 Context context, int appWidgetId, AppWidgetProviderInfo appWidget) {
-            return new QsbWidgetHostView(context);
+            return mViewFactory.newView(context);
         }
     }
+
+    public interface WidgetViewFactory {
+
+        QsbWidgetHostView newView(Context context);
+    }
 }
diff --git a/src/com/android/launcher3/qsb/QsbWidgetHostView.java b/src/com/android/launcher3/qsb/QsbWidgetHostView.java
index 7d8a4db..407812d 100644
--- a/src/com/android/launcher3/qsb/QsbWidgetHostView.java
+++ b/src/com/android/launcher3/qsb/QsbWidgetHostView.java
@@ -58,13 +58,9 @@
         try {
             super.onLayout(changed, left, top, right, bottom);
         } catch (final RuntimeException e) {
-            post(new Runnable() {
-                @Override
-                public void run() {
-                    // Update the widget with 0 Layout id, to reset the view to error view.
-                    updateAppWidget(new RemoteViews(getAppWidgetInfo().provider.getPackageName(), 0));
-                }
-            });
+            // Update the widget with 0 Layout id, to reset the view to error view.
+            post(() -> updateAppWidget(
+                    new RemoteViews(getAppWidgetInfo().provider.getPackageName(), 0)));
         }
     }
 
@@ -76,24 +72,16 @@
     @Override
     protected View getDefaultView() {
         View v = super.getDefaultView();
-        v.setOnClickListener(new OnClickListener() {
-            @Override
-            public void onClick(View view) {
-                Launcher.getLauncher(getContext()).startSearch("", false, null, true);
-            }
-        });
+        v.setOnClickListener((v2) ->
+                Launcher.getLauncher(getContext()).startSearch("", false, null, true));
         return v;
     }
 
     public static View getDefaultView(ViewGroup parent) {
         View v = LayoutInflater.from(parent.getContext())
                 .inflate(R.layout.qsb_default_view, parent, false);
-        v.findViewById(R.id.btn_qsb_search).setOnClickListener(new OnClickListener() {
-            @Override
-            public void onClick(View view) {
-                Launcher.getLauncher(view.getContext()).startSearch("", false, null, true);
-            }
-        });
+        v.findViewById(R.id.btn_qsb_search).setOnClickListener((v2) ->
+                Launcher.getLauncher(v2.getContext()).startSearch("", false, null, true));
         return v;
     }
 }