diff --git a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
index cce4c5b..9fdd21a 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddConfigWidgetTest.java
@@ -82,7 +82,6 @@
         runTest(false);
     }
 
-
     /**
      * @param acceptConfig accept the config activity
      */
@@ -101,7 +100,7 @@
         // Verify that the widget id is valid and bound
         assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
 
-        setResult(acceptConfig);
+        setResultAndWaitForAnimation(acceptConfig);
         if (acceptConfig) {
             Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
             assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
@@ -112,12 +111,20 @@
         }
     }
 
-    private void setResult(boolean success) {
-        mLauncher.executeAndWaitForWallpaperAnimation(() ->
-                        getInstrumentation().getTargetContext().sendBroadcast(
-                                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
-                                        success ? "clickOK" : "clickCancel")),
-                "setting widget coinfig result");
+    private static void setResult(boolean success) {
+        getInstrumentation().getTargetContext().sendBroadcast(
+                WidgetConfigActivity.getCommandIntent(WidgetConfigActivity.class,
+                        success ? "clickOK" : "clickCancel"));
+    }
+
+    private void setResultAndWaitForAnimation(boolean success) {
+        if (mLauncher.isLauncher3()) {
+            setResult(success);
+        } else {
+            mLauncher.executeAndWaitForWallpaperAnimation(
+                    () -> setResult(success),
+                    "setting widget coinfig result");
+        }
     }
 
     /**
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 0bc00a3..cbf1952 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -2180,7 +2180,8 @@
         };
     }
 
-    boolean isLauncher3() {
+    /** Returns whether the Launcher is a Launcher3 one */
+    public boolean isLauncher3() {
         if (mIsLauncher3 == null) {
             mIsLauncher3 = "com.android.launcher3".equals(getLauncherPackageName());
         }