diff --git a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
index dbf4b3e..17d0ac1 100644
--- a/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
+++ b/robolectric_tests/src/com/android/launcher3/testing/TestActivity.java
@@ -27,7 +27,12 @@
 
     @Override
     public BaseDragLayer getDragLayer() {
-        return null;
+        return new BaseDragLayer(this, /* attrs= */ null, /* alphaChannelCount= */ 1) {
+            @Override
+            public void recreateControllers() {
+                // Do nothing.
+            }
+        };
     }
 
     @Override
diff --git a/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java b/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
index 4e6f17c..a057a84 100644
--- a/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
+++ b/robolectric_tests/src/com/android/launcher3/widget/picker/search/WidgetsSearchBarControllerTest.java
@@ -19,24 +19,27 @@
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
-import android.content.Context;
 import android.view.View;
 import android.widget.ImageButton;
 
 import com.android.launcher3.ExtendedEditText;
 import com.android.launcher3.search.SearchAlgorithm;
+import com.android.launcher3.testing.TestActivity;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
+import org.robolectric.android.controller.ActivityController;
 
 import java.util.ArrayList;
 
@@ -44,7 +47,9 @@
 public class WidgetsSearchBarControllerTest {
 
     private WidgetsSearchBarController mController;
-    private Context mContext;
+    // TODO: Replace ActivityController with ActivityScenario, which is the recommended way for
+    // activity testing.
+    private ActivityController<TestActivity> mActivityController;
     private ExtendedEditText mEditText;
     private ImageButton mCancelButton;
     @Mock
@@ -55,13 +60,20 @@
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
-        mEditText = new ExtendedEditText(mContext);
-        mCancelButton = new ImageButton(mContext);
+        mActivityController = Robolectric.buildActivity(TestActivity.class);
+        TestActivity testActivity = mActivityController.setup().get();
+
+        mEditText = new ExtendedEditText(testActivity);
+        mCancelButton = new ImageButton(testActivity);
         mController = new WidgetsSearchBarController(
                 mSearchAlgorithm, mEditText, mCancelButton, mSearchModeListener);
     }
 
+    @After
+    public void tearDown() {
+        mActivityController.destroy();
+    }
+
     @Test
     public void onSearchResult_shouldInformSearchModeListener() {
         ArrayList<WidgetsListBaseEntry> entries = new ArrayList<>();
@@ -119,14 +131,18 @@
     public void cancelSearch_shouldInformSearchModeListenerToClearResultsAndExitSearch() {
         mCancelButton.performClick();
 
-        verify(mSearchModeListener).exitSearchMode();
+        // 1 time explicitly from the cancel button on click listener.
+        // Another from the setText("") the cancel button on click listener causing afterTextChange.
+        verify(mSearchModeListener, times(2)).exitSearchMode();
     }
 
     @Test
     public void cancelSearch_shouldCancelSearch() {
         mCancelButton.performClick();
 
-        verify(mSearchAlgorithm).cancel(true);
+        // 1 time explicitly from the cancel button on click listener.
+        // Another from the setText("") the cancel button on click listener causing afterTextChange.
+        verify(mSearchAlgorithm, times(2)).cancel(true);
         verifyNoMoreInteractions(mSearchAlgorithm);
     }