diff --git a/build.gradle b/build.gradle
index 3a812a9..51ac5a1 100644
--- a/build.gradle
+++ b/build.gradle
@@ -46,12 +46,16 @@
         androidTest {
             java.srcDirs = ['tests/src']
             res.srcDirs = ['tests/res']
-            manifest.srcFile "tests/AndroidManifest.xml"
+            manifest.srcFile "tests/AndroidManifest-common.xml"
         }
 
         aosp {
             manifest.srcFile "AndroidManifest.xml"
         }
+
+        aospAndroidTest {
+            manifest.srcFile "tests/AndroidManifest.xml"
+        }
     }
 }
 
diff --git a/src/com/android/launcher3/FolderInfo.java b/src/com/android/launcher3/FolderInfo.java
index 4c88e7e..5fff2e7 100644
--- a/src/com/android/launcher3/FolderInfo.java
+++ b/src/com/android/launcher3/FolderInfo.java
@@ -93,7 +93,7 @@
     }
 
     @Override
-    void onAddToDatabase(ContentWriter writer) {
+    public void onAddToDatabase(ContentWriter writer) {
         super.onAddToDatabase(writer);
         writer.put(LauncherSettings.Favorites.TITLE, title)
                 .put(LauncherSettings.Favorites.OPTIONS, options);
diff --git a/src/com/android/launcher3/ItemInfo.java b/src/com/android/launcher3/ItemInfo.java
index 82c7ab8..aec6c7d 100644
--- a/src/com/android/launcher3/ItemInfo.java
+++ b/src/com/android/launcher3/ItemInfo.java
@@ -166,7 +166,7 @@
     /**
      * Write the fields of this item to the DB
      */
-    void onAddToDatabase(ContentWriter writer) {
+    public void onAddToDatabase(ContentWriter writer) {
         if (screenId == Workspace.EXTRA_EMPTY_SCREEN_ID) {
             // We should never persist an item on the extra empty screen.
             throw new RuntimeException("Screen id should not be EXTRA_EMPTY_SCREEN_ID");
diff --git a/src/com/android/launcher3/LauncherAppWidgetInfo.java b/src/com/android/launcher3/LauncherAppWidgetInfo.java
index 285a93c..1e0f285 100644
--- a/src/com/android/launcher3/LauncherAppWidgetInfo.java
+++ b/src/com/android/launcher3/LauncherAppWidgetInfo.java
@@ -65,7 +65,7 @@
     /**
      * Indicates that the widget hasn't been instantiated yet.
      */
-    static final int NO_ID = -1;
+    public static final int NO_ID = -1;
 
     /**
      * Indicates that this is a locally defined widget and hence has no system allocated id.
@@ -126,7 +126,7 @@
     }
 
     @Override
-    void onAddToDatabase(ContentWriter writer) {
+    public void onAddToDatabase(ContentWriter writer) {
         super.onAddToDatabase(writer);
         writer.put(LauncherSettings.Favorites.APPWIDGET_ID, appWidgetId)
                 .put(LauncherSettings.Favorites.APPWIDGET_PROVIDER, providerName.flattenToString())
diff --git a/src/com/android/launcher3/MainThreadExecutor.java b/src/com/android/launcher3/MainThreadExecutor.java
index 866b17c..4ca0a59 100644
--- a/src/com/android/launcher3/MainThreadExecutor.java
+++ b/src/com/android/launcher3/MainThreadExecutor.java
@@ -16,65 +16,18 @@
 
 package com.android.launcher3;
 
-import android.os.Handler;
 import android.os.Looper;
 
-import java.util.List;
-import java.util.concurrent.AbstractExecutorService;
-import java.util.concurrent.TimeUnit;
+import com.android.launcher3.util.LooperExecuter;
 
 /**
  * An executor service that executes its tasks on the main thread.
  *
  * Shutting down this executor is not supported.
  */
-public class MainThreadExecutor extends AbstractExecutorService {
+public class MainThreadExecutor extends LooperExecuter {
 
-    private Handler mHandler = new Handler(Looper.getMainLooper());
-
-    @Override
-    public void execute(Runnable runnable) {
-        if (Looper.getMainLooper() == Looper.myLooper()) {
-            runnable.run();
-        } else {
-            mHandler.post(runnable);
-        }
-    }
-
-    /**
-     * Not supported and throws an exception when used.
-     */
-    @Override
-    @Deprecated
-    public void shutdown() {
-        throw new UnsupportedOperationException();
-    }
-
-    /**
-     * Not supported and throws an exception when used.
-     */
-    @Override
-    @Deprecated
-    public List<Runnable> shutdownNow() {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public boolean isShutdown() {
-        return false;
-    }
-
-    @Override
-    public boolean isTerminated() {
-        return false;
-    }
-
-    /**
-     * Not supported and throws an exception when used.
-     */
-    @Override
-    @Deprecated
-    public boolean awaitTermination(long l, TimeUnit timeUnit) throws InterruptedException {
-        throw new UnsupportedOperationException();
+    public MainThreadExecutor() {
+        super(Looper.getMainLooper());
     }
 }
diff --git a/src/com/android/launcher3/ShortcutInfo.java b/src/com/android/launcher3/ShortcutInfo.java
index 1a42395..b35dcb7 100644
--- a/src/com/android/launcher3/ShortcutInfo.java
+++ b/src/com/android/launcher3/ShortcutInfo.java
@@ -161,7 +161,7 @@
     }
 
     @Override
-    void onAddToDatabase(ContentWriter writer) {
+    public void onAddToDatabase(ContentWriter writer) {
         super.onAddToDatabase(writer);
         writer.put(LauncherSettings.BaseLauncherColumns.TITLE, title)
                 .put(LauncherSettings.BaseLauncherColumns.INTENT, getIntent())
diff --git a/src/com/android/launcher3/util/LooperExecuter.java b/src/com/android/launcher3/util/LooperExecuter.java
new file mode 100644
index 0000000..4db999b
--- /dev/null
+++ b/src/com/android/launcher3/util/LooperExecuter.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2017 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.util;
+
+import android.os.Handler;
+import android.os.Looper;
+
+import java.util.List;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Extension of {@link AbstractExecutorService} which executed on a provided looper.
+ */
+public class LooperExecuter extends AbstractExecutorService {
+
+    private final Handler mHandler;
+
+    public LooperExecuter(Looper looper) {
+        mHandler = new Handler(looper);
+    }
+
+    @Override
+    public void execute(Runnable runnable) {
+        if (mHandler.getLooper() == Looper.myLooper()) {
+            runnable.run();
+        } else {
+            mHandler.post(runnable);
+        }
+    }
+
+    /**
+     * Not supported and throws an exception when used.
+     */
+    @Override
+    @Deprecated
+    public void shutdown() {
+        throw new UnsupportedOperationException();
+    }
+
+    /**
+     * Not supported and throws an exception when used.
+     */
+    @Override
+    @Deprecated
+    public List<Runnable> shutdownNow() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    public boolean isShutdown() {
+        return false;
+    }
+
+    @Override
+    public boolean isTerminated() {
+        return false;
+    }
+
+    /**
+     * Not supported and throws an exception when used.
+     */
+    @Override
+    @Deprecated
+    public boolean awaitTermination(long l, TimeUnit timeUnit) throws InterruptedException {
+        throw new UnsupportedOperationException();
+    }
+}
diff --git a/tests/Android.mk b/tests/Android.mk
index 5103ced..e8797a7 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -20,6 +20,7 @@
 LOCAL_STATIC_JAVA_LIBRARIES := android-support-test ub-uiautomator mockito-target-minus-junit4
 
 LOCAL_SRC_FILES := $(call all-java-files-under, src)
+LOCAL_FULL_LIBS_MANIFEST_FILES := $(LOCAL_PATH)/AndroidManifest-common.xml
 
 LOCAL_SDK_VERSION := current
 LOCAL_MIN_SDK_VERSION := 21
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
new file mode 100644
index 0000000..763481a
--- /dev/null
+++ b/tests/AndroidManifest-common.xml
@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2017 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.
+-->
+
+<manifest
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.launcher3.tests">
+
+    <application android:debuggable="true">
+        <uses-library android:name="android.test.runner" />
+
+        <receiver android:name="com.android.launcher3.testcomponent.AppWidgetNoConfig">
+            <intent-filter>
+                <action android:name="android.appwidget.action.APPWIDGET_UPDATE" />
+            </intent-filter>
+            <meta-data android:name="android.appwidget.provider"
+                       android:resource="@xml/appwidget_no_config" />
+        </receiver>
+
+        <receiver android:name="com.android.launcher3.testcomponent.AppWidgetWithConfig">
+            <intent-filter>
+                <action android:name="android.appwidget.action.APPWIDGET_UPDATE" />
+            </intent-filter>
+            <meta-data android:name="android.appwidget.provider"
+                       android:resource="@xml/appwidget_with_config" />
+        </receiver>
+
+        <activity
+            android:name="com.android.launcher3.testcomponent.WidgetConfigActivity">
+            <intent-filter>
+                <action android:name="android.appwidget.action.APPWIDGET_CONFIGURE"/>
+            </intent-filter>
+        </activity>
+    </application>
+</manifest>
diff --git a/tests/res/layout/test_layout_appwidget_blue.xml b/tests/res/layout/test_layout_appwidget_blue.xml
new file mode 100644
index 0000000..8111978
--- /dev/null
+++ b/tests/res/layout/test_layout_appwidget_blue.xml
@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:orientation="vertical"
+    android:background="#FF0000FF"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent" />
\ No newline at end of file
diff --git a/tests/res/layout/test_layout_appwidget_red.xml b/tests/res/layout/test_layout_appwidget_red.xml
new file mode 100644
index 0000000..48d3e81
--- /dev/null
+++ b/tests/res/layout/test_layout_appwidget_red.xml
@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:orientation="vertical"
+    android:background="#FFFF0000"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent" />
\ No newline at end of file
diff --git a/tests/res/xml/appwidget_no_config.xml b/tests/res/xml/appwidget_no_config.xml
new file mode 100644
index 0000000..d24dfe3
--- /dev/null
+++ b/tests/res/xml/appwidget_no_config.xml
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="utf-8"?>
+<appwidget-provider
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:minWidth="180dp"
+    android:minHeight="110dp"
+    android:updatePeriodMillis="86400000"
+    android:initialLayout="@layout/test_layout_appwidget_red"
+    android:resizeMode="horizontal|vertical"
+    android:widgetCategory="home_screen">
+</appwidget-provider>
\ No newline at end of file
diff --git a/tests/res/xml/appwidget_with_config.xml b/tests/res/xml/appwidget_with_config.xml
new file mode 100644
index 0000000..3e96c6f
--- /dev/null
+++ b/tests/res/xml/appwidget_with_config.xml
@@ -0,0 +1,11 @@
+<?xml version="1.0" encoding="utf-8"?>
+<appwidget-provider
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:minWidth="180dp"
+    android:minHeight="110dp"
+    android:updatePeriodMillis="86400000"
+    android:initialLayout="@layout/test_layout_appwidget_blue"
+    android:configure="com.android.launcher3.testcomponent.WidgetConfigActivity"
+    android:resizeMode="horizontal|vertical"
+    android:widgetCategory="home_screen">
+</appwidget-provider>
\ No newline at end of file
diff --git a/tests/src/com/android/launcher3/testcomponent/AppWidgetNoConfig.java b/tests/src/com/android/launcher3/testcomponent/AppWidgetNoConfig.java
new file mode 100644
index 0000000..9b320d8
--- /dev/null
+++ b/tests/src/com/android/launcher3/testcomponent/AppWidgetNoConfig.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2017 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.testcomponent;
+
+import android.appwidget.AppWidgetProvider;
+
+/**
+ * A simple app widget without any configuration screen.
+ */
+public class AppWidgetNoConfig extends AppWidgetProvider {
+
+
+}
diff --git a/tests/src/com/android/launcher3/testcomponent/AppWidgetWithConfig.java b/tests/src/com/android/launcher3/testcomponent/AppWidgetWithConfig.java
new file mode 100644
index 0000000..033e6e6
--- /dev/null
+++ b/tests/src/com/android/launcher3/testcomponent/AppWidgetWithConfig.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2017 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.testcomponent;
+
+/**
+ * A simple app widget with configuration sceen.
+ */
+public class AppWidgetWithConfig extends AppWidgetNoConfig {
+
+}
diff --git a/tests/src/com/android/launcher3/testcomponent/WidgetConfigActivity.java b/tests/src/com/android/launcher3/testcomponent/WidgetConfigActivity.java
new file mode 100644
index 0000000..c0509bc
--- /dev/null
+++ b/tests/src/com/android/launcher3/testcomponent/WidgetConfigActivity.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2017 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.testcomponent;
+
+import android.app.Activity;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.os.Bundle;
+
+/**
+ * Simple activity for widget configuration
+ */
+public class WidgetConfigActivity extends Activity {
+
+    public static final String SUFFIX_FINISH = "-finish";
+    public static final String EXTRA_CODE = "code";
+    public static final String EXTRA_INTENT = "intent";
+
+    private final BroadcastReceiver mFinishReceiver = new BroadcastReceiver() {
+
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            WidgetConfigActivity.this.setResult(
+                    intent.getIntExtra(EXTRA_CODE, RESULT_CANCELED),
+                    (Intent) intent.getParcelableExtra(EXTRA_INTENT));
+            WidgetConfigActivity.this.finish();
+        }
+    };
+
+    private final String mAction = this.getClass().getName();
+
+    @Override
+    protected void onCreate(Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+        registerReceiver(mFinishReceiver, new IntentFilter(mAction + SUFFIX_FINISH));
+    }
+
+    @Override
+    protected void onResume() {
+        super.onResume();
+        sendBroadcast(new Intent(mAction).putExtra(Intent.EXTRA_INTENT, getIntent()));
+    }
+
+    @Override
+    protected void onDestroy() {
+        unregisterReceiver(mFinishReceiver);
+        super.onDestroy();
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java b/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java
index 42c6cd7..4bc40c6 100644
--- a/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java
+++ b/tests/src/com/android/launcher3/ui/LauncherInstrumentationTestCase.java
@@ -1,13 +1,27 @@
+/*
+ * Copyright (C) 2017 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.ui;
 
-import android.app.SearchManager;
-import android.appwidget.AppWidgetProviderInfo;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.graphics.Point;
 import android.os.ParcelFileDescriptor;
+import android.os.Process;
 import android.os.RemoteException;
 import android.os.SystemClock;
 import android.support.test.uiautomator.By;
@@ -19,22 +33,23 @@
 import android.test.InstrumentationTestCase;
 import android.view.MotionEvent;
 
-import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.MainThreadExecutor;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.testcomponent.AppWidgetNoConfig;
+import com.android.launcher3.testcomponent.AppWidgetWithConfig;
 import com.android.launcher3.util.ManagedProfileHeuristic;
 
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.util.Locale;
 import java.util.concurrent.Callable;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Base class for all instrumentation tests providing various utility methods.
@@ -42,6 +57,7 @@
 public class LauncherInstrumentationTestCase extends InstrumentationTestCase {
 
     public static final long DEFAULT_UI_TIMEOUT = 3000;
+    public static final long DEFAULT_WORKER_TIMEOUT_SECS = 5;
 
     protected UiDevice mDevice;
     protected Context mTargetContext;
@@ -233,18 +249,11 @@
      * Runs the callback on the UI thread and returns the result.
      */
     protected <T> T getOnUiThread(final Callable<T> callback) {
-        final AtomicReference<T> result = new AtomicReference<>(null);
         try {
-            runTestOnUiThread(new Runnable() {
-                @Override
-                public void run() {
-                    try {
-                        result.set(callback.call());
-                    } catch (Exception e) { }
-                }
-            });
-        } catch (Throwable t) { }
-        return result.get();
+            return new MainThreadExecutor().submit(callback).get();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
     }
 
     /**
@@ -252,35 +261,14 @@
      * @param hasConfigureScreen if true, a provider with a config screen is returned.
      */
     protected LauncherAppWidgetProviderInfo findWidgetProvider(final boolean hasConfigureScreen) {
-        LauncherAppWidgetProviderInfo info = getOnUiThread(new Callable<LauncherAppWidgetProviderInfo>() {
+        LauncherAppWidgetProviderInfo info =
+                getOnUiThread(new Callable<LauncherAppWidgetProviderInfo>() {
             @Override
             public LauncherAppWidgetProviderInfo call() throws Exception {
-                InvariantDeviceProfile idv = LauncherAppState.getIDP(mTargetContext);
-
-                ComponentName searchComponent = ((SearchManager) mTargetContext
-                        .getSystemService(Context.SEARCH_SERVICE)).getGlobalSearchActivity();
-                String searchPackage = searchComponent == null
-                        ? null : searchComponent.getPackageName();
-
-                for (AppWidgetProviderInfo info :
-                        AppWidgetManagerCompat.getInstance(mTargetContext).getAllProviders()) {
-                    if ((info.configure != null) ^ hasConfigureScreen) {
-                        continue;
-                    }
-                    // Exclude the widgets in search package, as Launcher already binds them in
-                    // QSB, so they can cause conflicts.
-                    if (info.provider.getPackageName().equals(searchPackage)) {
-                        continue;
-                    }
-                    LauncherAppWidgetProviderInfo widgetInfo = LauncherAppWidgetProviderInfo
-                            .fromProviderInfo(mTargetContext, info);
-                    if (widgetInfo.minSpanX >= idv.numColumns
-                            || widgetInfo.minSpanY >= idv.numRows) {
-                        continue;
-                    }
-                    return widgetInfo;
-                }
-                return null;
+                ComponentName cn = new ComponentName(getInstrumentation().getContext(),
+                        hasConfigureScreen ? AppWidgetWithConfig.class : AppWidgetNoConfig.class);
+                return AppWidgetManagerCompat.getInstance(mTargetContext)
+                        .findProvider(cn, Process.myUserHandle());
             }
         });
         if (info == null) {
diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
new file mode 100644
index 0000000..7cbd292
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -0,0 +1,215 @@
+/*
+ * Copyright (C) 2017 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.ui.widget;
+
+import android.app.Activity;
+import android.app.Application;
+import android.appwidget.AppWidgetManager;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.support.test.uiautomator.By;
+import android.support.test.uiautomator.UiObject2;
+import android.test.suitebuilder.annotation.LargeTest;
+import android.view.View;
+
+import com.android.launcher3.ItemInfo;
+import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherAppWidgetInfo;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.MainThreadExecutor;
+import com.android.launcher3.Workspace;
+import com.android.launcher3.testcomponent.WidgetConfigActivity;
+import com.android.launcher3.ui.LauncherInstrumentationTestCase;
+import com.android.launcher3.util.Condition;
+import com.android.launcher3.util.SimpleActivityMonitor;
+import com.android.launcher3.util.Wait;
+import com.android.launcher3.widget.WidgetCell;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Test to verify widget configuration is properly shown.
+ */
+@LargeTest
+public class AddConfigWidgetTest extends LauncherInstrumentationTestCase {
+
+    public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
+    public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 5;
+
+    private LauncherAppWidgetProviderInfo mWidgetInfo;
+    private SimpleActivityMonitor mActivityMonitor;
+    private MainThreadExecutor mMainThreadExecutor;
+    private AppWidgetManager mAppWidgetManager;
+
+    private int mWidgetId;
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        mWidgetInfo = findWidgetProvider(true /* hasConfigureScreen */);
+        mActivityMonitor = new SimpleActivityMonitor();
+        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
+                .registerActivityLifecycleCallbacks(mActivityMonitor);
+        mMainThreadExecutor = new MainThreadExecutor();
+        mAppWidgetManager = AppWidgetManager.getInstance(mTargetContext);
+    }
+
+    @Override
+    protected void tearDown() throws Exception {
+        ((Application) getInstrumentation().getTargetContext().getApplicationContext())
+                .unregisterActivityLifecycleCallbacks(mActivityMonitor);
+        super.tearDown();
+    }
+
+    public void testWidgetConfig() throws Throwable {
+        runTest(false, true);
+    }
+
+    public void testWidgetConfig_rotate() throws Throwable {
+        runTest(true, true);
+    }
+
+    public void testConfigCancelled() throws Throwable {
+        runTest(false, false);
+    }
+
+    public void testConfigCancelled_rotate() throws Throwable {
+        runTest(true, false);
+    }
+
+    /**
+     * @param rotateConfig should the config screen be rotated
+     * @param acceptConfig accept the config activity
+     */
+    private void runTest(boolean rotateConfig, boolean acceptConfig) throws Throwable {
+        lockRotation(true);
+
+        clearHomescreen();
+        startLauncher();
+
+        // Open widget tray and wait for load complete.
+        final UiObject2 widgetContainer = openWidgetsTray();
+        assertTrue(Wait.atMost(Condition.minChildCount(widgetContainer, 2), DEFAULT_UI_TIMEOUT));
+
+        // Drag widget to homescreen
+        WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
+        UiObject2 widget = scrollAndFind(widgetContainer, By.clazz(WidgetCell.class)
+                .hasDescendant(By.text(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))));
+        dragToWorkspace(widget, false);
+        // Widget id for which the config activity was opened
+        mWidgetId = monitor.getWidgetId();
+
+        if (rotateConfig) {
+            // Rotate the screen and verify that the config activity is recreated
+            monitor = new WidgetConfigStartupMonitor();
+            lockRotation(false);
+            assertEquals(mWidgetId, monitor.getWidgetId());
+        }
+
+        // Verify that the widget id is valid and bound
+        assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
+
+        if (acceptConfig) {
+            setResult(Activity.RESULT_OK);
+            assertTrue(Wait.atMost(new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT));
+            assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
+        } else {
+            setResult(Activity.RESULT_CANCELED);
+            // Verify that the widget id is deleted.
+            assertTrue(Wait.atMost(new Condition() {
+                @Override
+                public boolean isTrue() throws Throwable {
+                    return mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null;
+                }
+            }, DEFAULT_ACTIVITY_TIMEOUT));
+        }
+    }
+
+    private void setResult(int resultCode) {
+        String action = WidgetConfigActivity.class.getName() + WidgetConfigActivity.SUFFIX_FINISH;
+        getInstrumentation().getTargetContext().sendBroadcast(
+                new Intent(action).putExtra(WidgetConfigActivity.EXTRA_CODE, resultCode));
+    }
+
+    /**
+     * Condition for searching widget id
+     */
+    private class WidgetSearchCondition extends Condition
+            implements Callable<Boolean>, Workspace.ItemOperator {
+
+        @Override
+        public boolean isTrue() throws Throwable {
+            return mMainThreadExecutor.submit(this).get();
+        }
+
+        @Override
+        public boolean evaluate(ItemInfo info, View view) {
+            return info instanceof LauncherAppWidgetInfo &&
+                    ((LauncherAppWidgetInfo) info).providerName.equals(mWidgetInfo.provider) &&
+                    ((LauncherAppWidgetInfo) info).appWidgetId == mWidgetId;
+        }
+
+        @Override
+        public Boolean call() throws Exception {
+            // Find the resumed launcher
+            Launcher launcher = null;
+            for (Activity a : mActivityMonitor.resumed) {
+                if (a instanceof Launcher) {
+                    launcher = (Launcher) a;
+                }
+            }
+            if (launcher == null) {
+                return false;
+            }
+            return launcher.getWorkspace().getFirstMatch(this) != null;
+        }
+    }
+
+    /**
+     * Broadcast receiver for receiving widget config activity status.
+     */
+    private class WidgetConfigStartupMonitor extends BroadcastReceiver {
+
+        private final CountDownLatch latch = new CountDownLatch(1);
+        private Intent mIntent;
+
+        WidgetConfigStartupMonitor() {
+            getInstrumentation().getTargetContext().registerReceiver(this,
+                    new IntentFilter(WidgetConfigActivity.class.getName()));
+        }
+
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            mIntent = intent.getParcelableExtra(Intent.EXTRA_INTENT);
+            latch.countDown();
+        }
+
+        public int getWidgetId() throws InterruptedException {
+            latch.await(DEFAULT_BROADCAST_TIMEOUT_SECS, TimeUnit.SECONDS);
+            getInstrumentation().getTargetContext().unregisterReceiver(this);
+            assertNotNull(mIntent);
+            assertEquals(AppWidgetManager.ACTION_APPWIDGET_CONFIGURE, mIntent.getAction());
+            int widgetId = mIntent.getIntExtra(AppWidgetManager.EXTRA_APPWIDGET_ID,
+                    LauncherAppWidgetInfo.NO_ID);
+            assertNotSame(widgetId, LauncherAppWidgetInfo.NO_ID);
+            return widgetId;
+        }
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
similarity index 74%
rename from tests/src/com/android/launcher3/ui/AddWidgetTest.java
rename to tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index d536af2..b7e1ca9 100644
--- a/tests/src/com/android/launcher3/ui/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -1,4 +1,19 @@
-package com.android.launcher3.ui;
+/*
+ * Copyright (C) 2017 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.ui.widget;
 
 import android.support.test.uiautomator.By;
 import android.support.test.uiautomator.UiObject2;
@@ -10,6 +25,7 @@
 import com.android.launcher3.LauncherAppWidgetInfo;
 import com.android.launcher3.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.Workspace.ItemOperator;
+import com.android.launcher3.ui.LauncherInstrumentationTestCase;
 import com.android.launcher3.util.Condition;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.widget.WidgetCell;
diff --git a/tests/src/com/android/launcher3/BindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
similarity index 90%
rename from tests/src/com/android/launcher3/BindWidgetTest.java
rename to tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
index 575b42b..df2b662 100644
--- a/tests/src/com/android/launcher3/BindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/BindWidgetTest.java
@@ -1,4 +1,19 @@
-package com.android.launcher3;
+/*
+ * Copyright (C) 2017 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.ui.widget;
 
 import android.appwidget.AppWidgetHost;
 import android.content.ComponentName;
@@ -12,10 +27,19 @@
 import android.support.test.uiautomator.UiSelector;
 import android.test.suitebuilder.annotation.LargeTest;
 
+import com.android.launcher3.Launcher;
+import com.android.launcher3.LauncherAppWidgetHostView;
+import com.android.launcher3.LauncherAppWidgetInfo;
+import com.android.launcher3.LauncherAppWidgetProviderInfo;
+import com.android.launcher3.LauncherModel;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.PendingAppWidgetHostView;
+import com.android.launcher3.Workspace;
 import com.android.launcher3.compat.AppWidgetManagerCompat;
 import com.android.launcher3.compat.PackageInstallerCompat;
 import com.android.launcher3.ui.LauncherInstrumentationTestCase;
 import com.android.launcher3.util.ContentWriter;
+import com.android.launcher3.util.LooperExecuter;
 import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.WidgetHostViewLoader;
 
@@ -315,14 +339,11 @@
     /**
      * Blocks the current thread until all the jobs in the main worker thread are complete.
      */
-    private void waitUntilLoaderIdle() throws InterruptedException {
-        final CountDownLatch latch = new CountDownLatch(1);
-        LauncherModel.sWorker.post(new Runnable() {
-            @Override
-            public void run() {
-                latch.countDown();
-            }
-        });
-        assertTrue(latch.await(5, TimeUnit.SECONDS));
+    private void waitUntilLoaderIdle() throws Exception {
+        new LooperExecuter(LauncherModel.getWorkerLooper())
+                .submit(new Runnable() {
+                    @Override
+                    public void run() { }
+                }).get(DEFAULT_WORKER_TIMEOUT_SECS, TimeUnit.SECONDS);
     }
 }
diff --git a/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java b/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java
new file mode 100644
index 0000000..6154ab6
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/SimpleActivityMonitor.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2017 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.util;
+
+import android.app.Activity;
+import android.app.Application.*;
+import android.os.Bundle;
+
+import java.util.ArrayList;
+
+/**
+ * Simple monitor to keep a list of active activities.
+ */
+public class SimpleActivityMonitor implements ActivityLifecycleCallbacks {
+
+    public final ArrayList<Activity> created = new ArrayList<>();
+    public final ArrayList<Activity> started = new ArrayList<>();
+    public final ArrayList<Activity> resumed = new ArrayList<>();
+
+    @Override
+    public void onActivityCreated(Activity activity, Bundle bundle) {
+        created.add(activity);
+    }
+
+    @Override
+    public void onActivityStarted(Activity activity) {
+        started.add(activity);
+    }
+
+    @Override
+    public void onActivityResumed(Activity activity) {
+        resumed.add(activity);
+    }
+
+    @Override
+    public void onActivityPaused(Activity activity) {
+        resumed.remove(activity);
+    }
+
+    @Override
+    public void onActivityStopped(Activity activity) {
+        started.remove(activity);
+    }
+
+    @Override
+    public void onActivitySaveInstanceState(Activity activity, Bundle bundle) { }
+
+    @Override
+    public void onActivityDestroyed(Activity activity) {
+        created.remove(activity);
+    }
+}