diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 7e8890a..f81463a 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -29,6 +29,9 @@
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.util.Log;
+import android.view.MotionEvent;
+
 import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.BySelector;
@@ -36,8 +39,6 @@
 import androidx.test.uiautomator.UiDevice;
 import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
-import android.util.Log;
-import android.view.MotionEvent;
 
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
@@ -90,7 +91,7 @@
     @Before
     public void setUp() throws Exception {
         mDevice = UiDevice.getInstance(getInstrumentation());
-        mLauncher = new LauncherInstrumentation(mDevice);
+        mLauncher = new LauncherInstrumentation(getInstrumentation());
         mTargetContext = InstrumentationRegistry.getTargetContext();
         mTargetPackage = mTargetContext.getPackageName();
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 74a91a4..09bf403 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -23,16 +23,19 @@
 import static org.junit.Assert.fail;
 
 import android.app.ActivityManager;
+import android.app.Instrumentation;
 import android.app.UiAutomation;
 import android.os.Bundle;
 import android.provider.Settings;
-import androidx.test.InstrumentationRegistry;
+import android.view.accessibility.AccessibilityEvent;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiDevice;
 import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
-import android.view.accessibility.AccessibilityEvent;
 
 import com.android.launcher3.TestProtocol;
 import com.android.quickstep.SwipeUpSetting;
@@ -40,9 +43,6 @@
 import java.lang.ref.WeakReference;
 import java.util.concurrent.TimeoutException;
 
-import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
-
 /**
  * The main tapl object. The only object that can be explicitly constructed by the using code. It
  * produces all other objects.
@@ -90,17 +90,19 @@
     private final UiDevice mDevice;
     private final boolean mSwipeUpEnabled;
     private Boolean mSwipeUpEnabledOverride = null;
+    private final Instrumentation mInstrumentation;
 
     /**
      * Constructs the root of TAPL hierarchy. You get all other objects from it.
      */
-    public LauncherInstrumentation(UiDevice device) {
-        mDevice = device;
+    public LauncherInstrumentation(Instrumentation instrumentation) {
+        mInstrumentation = instrumentation;
+        mDevice = UiDevice.getInstance(instrumentation);
         final boolean swipeUpEnabledDefault =
                 !SwipeUpSetting.isSwipeUpSettingAvailable() ||
                         SwipeUpSetting.isSwipeUpEnabledDefaultValue();
         mSwipeUpEnabled = Settings.Secure.getInt(
-                InstrumentationRegistry.getTargetContext().getContentResolver(),
+                instrumentation.getTargetContext().getContentResolver(),
                 SWIPE_UP_SETTING_NAME,
                 swipeUpEnabledDefault ? 1 : 0) == 1;
         assertTrue("Device must run in a test harness", ActivityManager.isRunningInTestHarness());
@@ -162,9 +164,8 @@
             UiAutomation.AccessibilityEventFilter eventFilter, String message) {
         try {
             final AccessibilityEvent event =
-                    InstrumentationRegistry.getInstrumentation().getUiAutomation()
-                            .executeAndWaitForEvent(
-                                    command, eventFilter, WAIT_TIME_MS);
+                    mInstrumentation.getUiAutomation().executeAndWaitForEvent(
+                            command, eventFilter, WAIT_TIME_MS);
             assertNotNull("executeAndWaitForEvent returned null (this can't happen)", event);
             return (Bundle) event.getParcelableData();
         } catch (TimeoutException e) {