diff --git a/utils/dialogs/Android.mk b/utils/dialogs/Android.mk
index 9440bfd..8a625bb 100644
--- a/utils/dialogs/Android.mk
+++ b/utils/dialogs/Android.mk
@@ -20,7 +20,7 @@
 LOCAL_MODULE_TAGS := tests
 
 LOCAL_PACKAGE_NAME := DialogDismissalUtil
-LOCAL_STATIC_JAVA_LIBRARIES := first-party-app-helpers ub-uiautomator AuptLib
+LOCAL_STATIC_JAVA_LIBRARIES := first-party-app-helpers ub-uiautomator AuptLib android-support-test
 LOCAL_SRC_FILES := $(call all-java-files-under, src)
 LOCAL_SDK_VERSION := 23
 
diff --git a/utils/dialogs/src/com/android/dialogutils/DismissDialogsInstrumentation.java b/utils/dialogs/src/com/android/dialogutils/DismissDialogsInstrumentation.java
index 5a42255..8f6cdc0 100644
--- a/utils/dialogs/src/com/android/dialogutils/DismissDialogsInstrumentation.java
+++ b/utils/dialogs/src/com/android/dialogutils/DismissDialogsInstrumentation.java
@@ -22,30 +22,29 @@
 import android.os.Environment;
 import android.os.RemoteException;
 import android.os.SystemClock;
-import android.support.test.aupt.UiWatchers;
-import android.support.test.uiautomator.UiDevice;
-import android.support.test.uiautomator.Until;
-import android.util.Log;
-
-import android.platform.test.helpers.IStandardAppHelper;
 import android.platform.test.helpers.ChromeHelperImpl;
+import android.platform.test.helpers.GmailHelperImpl;
 import android.platform.test.helpers.GoogleCameraHelperImpl;
 import android.platform.test.helpers.GoogleKeyboardHelperImpl;
-import android.platform.test.helpers.GmailHelperImpl;
+import android.platform.test.helpers.IStandardAppHelper;
 import android.platform.test.helpers.MapsHelperImpl;
 import android.platform.test.helpers.PhotosHelperImpl;
 import android.platform.test.helpers.PlayMoviesHelperImpl;
 import android.platform.test.helpers.PlayMusicHelperImpl;
 import android.platform.test.helpers.PlayStoreHelperImpl;
 import android.platform.test.helpers.YouTubeHelperImpl;
+import android.support.test.InstrumentationRegistry;
+import android.support.test.aupt.UiWatchers;
 import android.support.test.launcherhelper.ILauncherStrategy;
 import android.support.test.launcherhelper.LauncherStrategyFactory;
-
+import android.support.test.uiautomator.UiDevice;
+import android.support.test.uiautomator.Until;
+import android.util.Log;
 import java.io.File;
 import java.io.IOException;
-import java.lang.NoSuchMethodException;
-import java.lang.InstantiationException;
 import java.lang.IllegalAccessException;
+import java.lang.InstantiationException;
+import java.lang.NoSuchMethodException;
 import java.lang.ReflectiveOperationException;
 import java.lang.reflect.InvocationTargetException;
 import java.util.HashMap;
@@ -124,6 +123,8 @@
     public void onStart() {
         super.onStart();
 
+        InstrumentationRegistry.registerInstance(this, new Bundle());
+
         mDevice = UiDevice.getInstance(this);
 
         UiWatchers watcherManager = new UiWatchers();