diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index d7f3547..7312dab 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -15,25 +15,27 @@
  */
 package com.android.launcher3.ui.widget;
 
+import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
+
+import static com.android.launcher3.util.LauncherBindableItemsContainer.ItemOperator;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
-import static org.junit.Assert.assertNull;
 
 import android.appwidget.AppWidgetManager;
 import android.content.Intent;
-import android.util.Log;
+import android.view.View;
 
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
-import com.android.launcher3.tapl.Widget;
-import com.android.launcher3.tapl.WidgetResizeFrame;
-import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.testcomponent.WidgetConfigActivity;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -50,7 +52,6 @@
 @RunWith(AndroidJUnit4.class)
 public class AddConfigWidgetTest extends AbstractLauncherUiTest {
 
-    private static final String WIDGET_PROVIDER_INFO = "b/276794291";
     @Rule
     public ShellCommandRule mGrantWidgetRule = ShellCommandRule.grantWidgetBind();
 
@@ -88,33 +89,51 @@
         clearHomescreen();
         mDevice.pressHome();
 
-        final Widgets widgets = mLauncher.getWorkspace().openAllWidgets();
-
         // Drag widget to homescreen
         WidgetConfigStartupMonitor monitor = new WidgetConfigStartupMonitor();
-        WidgetResizeFrame resizeFrame =
-                widgets.getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
-                        .dragConfigWidgetToWorkspace(acceptConfig);
+        mLauncher.getWorkspace()
+                .openAllWidgets()
+                .getWidget(mWidgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .dragToWorkspace(true, false);
         // Widget id for which the config activity was opened
         mWidgetId = monitor.getWidgetId();
 
-        Log.e(WIDGET_PROVIDER_INFO,
-                "InstalledProviders count: " + mAppWidgetManager.getInstalledProviders().size());
-
         // Verify that the widget id is valid and bound
         assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
 
+        setResult(acceptConfig);
         if (acceptConfig) {
-            assertNotNull("Widget resize frame not shown after widget added", resizeFrame);
-            resizeFrame.dismiss();
-
-            final Widget widget =
-                    mLauncher.getWorkspace().tryGetWidget(mWidgetInfo.label, DEFAULT_UI_TIMEOUT);
-            assertNotNull("Widget not found on the workspace", widget);
+            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+            assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
         } else {
-            final Widget widget =
-                    mLauncher.getWorkspace().tryGetWidget(mWidgetInfo.label, DEFAULT_UI_TIMEOUT);
-            assertNull("Widget unexpectedly found on the workspace", widget);
+            // Verify that the widget id is deleted.
+            Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
+                    DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+        }
+    }
+
+    private void setResult(boolean success) {
+        getInstrumentation().getTargetContext().sendBroadcast(
+                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
+                        success ? "clickOK" : "clickCancel"));
+    }
+
+    /**
+     * Condition for searching widget id
+     */
+    private class WidgetSearchCondition implements Wait.Condition, ItemOperator {
+
+        @Override
+        public boolean isTrue() throws Throwable {
+            return mMainThreadExecutor.submit(mActivityMonitor.itemExists(this)).get();
+        }
+
+        @Override
+        public boolean evaluate(ItemInfo info, View view) {
+            return info instanceof LauncherAppWidgetInfo
+                    && ((LauncherAppWidgetInfo) info).providerName.getClassName().equals(
+                            mWidgetInfo.provider.getClassName())
+                    && ((LauncherAppWidgetInfo) info).appWidgetId == mWidgetId;
         }
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/Widget.java b/tests/tapl/com/android/launcher3/tapl/Widget.java
index d440903..e7016a3 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widget.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widget.java
@@ -17,11 +17,7 @@
 package com.android.launcher3.tapl;
 
 import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-import androidx.test.uiautomator.By;
-import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiObject2;
-import androidx.test.uiautomator.Until;
 
 import com.android.launcher3.testing.shared.TestProtocol;
 
@@ -69,10 +65,7 @@
      */
     @NonNull
     public WidgetResizeFrame dragWidgetToWorkspace() {
-        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
-            return dragWidgetToWorkspace(/* configurable= */ false, /* acceptsConfig= */ false, -1,
-                    -1, 1, 1);
-        }
+        return dragWidgetToWorkspace(-1, -1, 1, 1);
     }
 
     /**
@@ -84,22 +77,19 @@
         try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
              LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
                      "Dragging widget to workspace cell " + cellX + "," + cellY)) {
-            return dragWidgetToWorkspace(/* configurable= */ false, /* acceptsConfig= */ false,
-                    cellX, cellY, spanX, spanY);
-        }
-    }
+            if (cellX == -1 || cellY == -1) {
+                internalDragToWorkspace(/* startsActivity= */ false, /* isWidgetShortcut= */
+                        false);
+            } else {
+                dragToWorkspaceCellPosition(/* startsActivity= */ false, /* isWidgetShortcut= */
+                        false, cellX, cellY, spanX, spanY);
+            }
 
-    /**
-     * Drags a configurable widget from the widgets container to the workspace, either accepts or
-     * cancels the configuration based on {@code acceptsConfig}, and returns the resize frame that
-     * is shown if the widget is added.
-     */
-    @Nullable
-    public WidgetResizeFrame dragConfigWidgetToWorkspace(boolean acceptsConfig) {
-        // TODO(b/239438337, fransebas) add correct event checking for this case
-        //try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck()) {
-        return dragWidgetToWorkspace(/* configurable= */ true, acceptsConfig, -1, -1, 1, 1);
-        //}
+            try (LauncherInstrumentation.Closable closable = mLauncher.addContextLayer(
+                    "want to get widget resize frame")) {
+                return new WidgetResizeFrame(mLauncher);
+            }
+        }
     }
 
     /**
@@ -122,53 +112,4 @@
                 isWidgetShortcut,
                 launchable::addExpectedEventsForLongClick);
     }
-
-    /**
-     * Drags a widget from the widgets container to the workspace and returns the resize frame that
-     * is shown after the widget is added.
-     *
-     * <p> If {@code configurable} is true, then either accepts or cancels the configuration based
-     * on {@code acceptsConfig}.
-     * <p> If either {@code cellX} or {@code cellY} are negative, then a default location would be
-     * chosen
-     *
-     * @param configurable  if the widget has a configuration activity.
-     * @param acceptsConfig if the widget has a configuration, then if we should accept it or
-     *                      cancel it
-     * @param cellX         X position to drop the widget in the workspace
-     * @param cellY         Y position to drop the widget in the workspace
-     * @return returns the given resize frame of the widget after being dropped, if
-     * configurable is true and acceptsConfig is false then the widget would not be places and will
-     * be cancel and it returns null.
-     */
-    @Nullable
-    private WidgetResizeFrame dragWidgetToWorkspace(boolean configurable, boolean acceptsConfig,
-            int cellX, int cellY, int spanX, int spanY) {
-        if (cellX == -1 || cellY == -1) {
-            internalDragToWorkspace(/* startsActivity= */ configurable, /* isWidgetShortcut= */
-                    false);
-        } else {
-            dragToWorkspaceCellPosition(/* startsActivity= */ configurable, /* isWidgetShortcut= */
-                    false, cellX, cellY, spanX, spanY);
-        }
-
-        if (configurable) {
-            // Configure the widget.
-            BySelector selector = By.text(acceptsConfig ? "OK" : "Cancel");
-            mLauncher.getDevice()
-                    .wait(Until.findObject(selector), LauncherInstrumentation.WAIT_TIME_MS)
-                    .click();
-
-            // If the widget configuration was cancelled, then the widget wasn't added to the home
-            // screen. In that case, we cannot return a resize frame.
-            if (!acceptsConfig) {
-                return null;
-            }
-        }
-
-        try (LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
-                "want to get widget resize frame")) {
-            return new WidgetResizeFrame(mLauncher);
-        }
-    }
 }