Adding some widget addition flow tests

> Added two dummy widget providers: with config and without config
> Added tests for verify widget config flow

Change-Id: I4577f085abe8f8b82047b644c71cc9065358153a

14 files changed, 554 insertions, 53 deletions

diff --git a/tests/Android.mk b/tests/Android.mk
index 5103ced7c..e8797a75c 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -20,6 +20,7 @@ LOCAL_MODULE_TAGS := tests
 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 000000000..763481ae2
--- /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 000000000..8111978c9
--- /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 000000000..48d3e8107
--- /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 000000000..d24dfe34d
--- /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 000000000..3e96c6f2f
--- /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 000000000..9b320d83b
--- /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 000000000..033e6e674
--- /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 000000000..c0509bc24
--- /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 42c6cd729..4bc40c68f 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.support.test.uiautomator.Until;
 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 @@ import java.util.concurrent.atomic.AtomicReference;
 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 @@ public class LauncherInstrumentationTestCase extends InstrumentationTestCase {
      * 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 @@ public class LauncherInstrumentationTestCase extends InstrumentationTestCase {
      * @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 000000000..7cbd29283
--- /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
index d536af27e..b7e1ca903 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.Launcher;
 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
index 575b42b87..df2b66285 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.os.Bundle;
 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 @@ public class BindWidgetTest extends LauncherInstrumentationTestCase {
     /**
      * 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 000000000..6154ab6c5
--- /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);
+    }
+}