Merge "Go to root  from child directory when root tapped. Bug: 27060001" into nyc-dev
diff --git a/res/layout/directory_cluster.xml b/res/layout/directory_cluster.xml
index 2fa09d3..d84ef08 100644
--- a/res/layout/directory_cluster.xml
+++ b/res/layout/directory_cluster.xml
@@ -27,6 +27,7 @@
 
     <FrameLayout
         android:id="@+id/container_directory"
+        android:clipToPadding="false"
         android:layout_width="match_parent"
         android:layout_height="0dp"
         android:layout_weight="1" />
diff --git a/res/layout/fragment_directory.xml b/res/layout/fragment_directory.xml
index 223d729..d0364ff 100644
--- a/res/layout/fragment_directory.xml
+++ b/res/layout/fragment_directory.xml
@@ -17,8 +17,10 @@
 <com.android.documentsui.DirectoryView xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
-    android:orientation="vertical"
-    android:animateLayoutChanges="true">
+    android:background="@color/window_background"
+    android:outlineProvider="bounds"
+    android:elevation="4dp"
+    android:orientation="vertical">
 
     <ProgressBar
         android:id="@+id/progressbar"
@@ -43,6 +45,7 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"
         android:orientation="vertical"
+        android:background="@color/window_background"
         android:visibility="gone">
 
         <LinearLayout
@@ -81,7 +84,6 @@
 
         <android.support.v7.widget.RecyclerView
             android:id="@+id/list"
-            android:background="@color/window_background"
             android:scrollbars="vertical"
             android:layout_width="match_parent"
             android:layout_height="match_parent"
diff --git a/src/com/android/documentsui/Events.java b/src/com/android/documentsui/Events.java
index 99b425e..14d4e2d 100644
--- a/src/com/android/documentsui/Events.java
+++ b/src/com/android/documentsui/Events.java
@@ -91,6 +91,8 @@
             case KeyEvent.KEYCODE_DPAD_RIGHT:
             case KeyEvent.KEYCODE_MOVE_HOME:
             case KeyEvent.KEYCODE_MOVE_END:
+            case KeyEvent.KEYCODE_PAGE_UP:
+            case KeyEvent.KEYCODE_PAGE_DOWN:
                 return true;
             default:
                 return false;
diff --git a/src/com/android/documentsui/dirlist/DirectoryFragment.java b/src/com/android/documentsui/dirlist/DirectoryFragment.java
index 70bee3c..b6c27b5 100644
--- a/src/com/android/documentsui/dirlist/DirectoryFragment.java
+++ b/src/com/android/documentsui/dirlist/DirectoryFragment.java
@@ -53,9 +53,7 @@
 import android.support.design.widget.Snackbar;
 import android.support.v7.widget.GridLayoutManager;
 import android.support.v7.widget.GridLayoutManager.SpanSizeLookup;
-import android.support.v7.widget.LinearLayoutManager;
 import android.support.v7.widget.RecyclerView;
-import android.support.v7.widget.RecyclerView.LayoutManager;
 import android.support.v7.widget.RecyclerView.OnItemTouchListener;
 import android.support.v7.widget.RecyclerView.RecyclerListener;
 import android.support.v7.widget.RecyclerView.ViewHolder;
@@ -94,11 +92,13 @@
 import com.android.documentsui.Shared;
 import com.android.documentsui.Snackbars;
 import com.android.documentsui.State;
+import com.android.documentsui.State.ViewMode;
 import com.android.documentsui.dirlist.MultiSelectManager.Selection;
 import com.android.documentsui.model.DocumentInfo;
 import com.android.documentsui.model.DocumentStack;
 import com.android.documentsui.model.RootInfo;
 import com.android.documentsui.services.FileOperationService;
+import com.android.documentsui.services.FileOperationService.OpType;
 import com.android.documentsui.services.FileOperations;
 import com.google.common.collect.Lists;
 
@@ -155,9 +155,7 @@
     private LoaderCallbacks<DirectoryResult> mCallbacks;
     private FragmentTuner mTuner;
     private DocumentClipper mClipper;
-    // These are lazily initialized.
-    private LinearLayoutManager mListLayout;
-    private GridLayoutManager mGridLayout;
+    private GridLayoutManager mLayout;
     private int mColumnCount = 1;  // This will get updated when layout changes.
 
     private MessageBar mMessageBar;
@@ -182,22 +180,6 @@
                     }
                 });
 
-        // TODO: Rather than update columns on layout changes, push this
-        // code (or something like it) into GridLayoutManager.
-        mRecView.addOnLayoutChangeListener(
-                new View.OnLayoutChangeListener() {
-
-                    @Override
-                    public void onLayoutChange(
-                            View v, int left, int top, int right, int bottom, int oldLeft,
-                            int oldTop, int oldRight, int oldBottom) {
-                        mColumnCount = calculateColumnCount();
-                        if (mGridLayout != null) {
-                            mGridLayout.setSpanCount(mColumnCount);
-                        }
-                    }
-                });
-
         mRecView.setItemAnimator(new DirectoryItemAnimator(getActivity()));
 
         // TODO: Add a divider between views (which might use RecyclerView.ItemDecoration).
@@ -240,6 +222,13 @@
 
         mRecView.setAdapter(mAdapter);
 
+        mLayout = new GridLayoutManager(getContext(), mColumnCount);
+        SpanSizeLookup lookup = mAdapter.createSpanSizeLookup();
+        if (lookup != null) {
+            mLayout.setSpanSizeLookup(lookup);
+        }
+        mRecView.setLayoutManager(mLayout);
+
         mGestureDetector = new ListeningGestureDetector(this.getContext(), new GestureListener());
 
         mRecView.addOnItemTouchListener(mGestureDetector);
@@ -432,41 +421,28 @@
     }
 
     /**
-     * Returns a {@code LayoutManager} for {@code mode}, lazily initializing
-     * classes as needed.
+     * Updates the layout after the view mode switches.
+     * @param mode The new view mode.
      */
-    private void updateLayout(int mode) {
-        final LayoutManager layout;
-        switch (mode) {
-            case MODE_GRID:
-                if (mGridLayout == null) {
-                    mGridLayout = new GridLayoutManager(getContext(), mColumnCount);
-                    SpanSizeLookup lookup = mAdapter.createSpanSizeLookup();
-                    if (lookup != null) {
-                        mGridLayout.setSpanSizeLookup(lookup);
-                    }
-                }
-                layout = mGridLayout;
-                break;
-            case MODE_LIST:
-                if (mListLayout == null) {
-                    mListLayout = new LinearLayoutManager(getContext());
-                }
-                layout = mListLayout;
-                break;
-            default:
-                throw new IllegalArgumentException("Unsupported layout mode: " + mode);
+    private void updateLayout(@ViewMode int mode) {
+        mColumnCount = calculateColumnCount(mode);
+        if (mLayout != null) {
+            mLayout.setSpanCount(mColumnCount);
         }
 
         int pad = getDirectoryPadding(mode);
         mRecView.setPadding(pad, pad, pad, pad);
-        // setting layout manager automatically invalidates existing ViewHolders.
-        mRecView.setLayoutManager(layout);
+        mRecView.requestLayout();
         mSelectionManager.handleLayoutChanged();  // RecyclerView doesn't do this for us
         mIconHelper.setViewMode(mode);
     }
 
-    private int calculateColumnCount() {
+    private int calculateColumnCount(@ViewMode int mode) {
+        if (mode == MODE_LIST) {
+            // List mode is a "grid" with 1 column.
+            return 1;
+        }
+
         int cellWidth = getResources().getDimensionPixelSize(R.dimen.grid_width);
         int cellMargin = 2 * getResources().getDimensionPixelSize(R.dimen.grid_item_margin);
         int viewPadding = mRecView.getPaddingLeft() + mRecView.getPaddingRight();
@@ -478,14 +454,12 @@
         return columnCount;
     }
 
-    private int getDirectoryPadding(int mode) {
+    private int getDirectoryPadding(@ViewMode int mode) {
         switch (mode) {
             case MODE_GRID:
-                return getResources().getDimensionPixelSize(
-                        R.dimen.grid_container_padding);
+                return getResources().getDimensionPixelSize(R.dimen.grid_container_padding);
             case MODE_LIST:
-                return getResources().getDimensionPixelSize(
-                        R.dimen.list_container_padding);
+                return getResources().getDimensionPixelSize(R.dimen.list_container_padding);
             default:
                 throw new IllegalArgumentException("Unsupported layout mode: " + mode);
         }
@@ -784,7 +758,7 @@
                 .show();
     }
 
-    private void transferDocuments(final Selection selected, final int mode) {
+    private void transferDocuments(final Selection selected, final @OpType int mode) {
         // Pop up a dialog to pick a destination.  This is inadequate but works for now.
         // TODO: Implement a picker that is to spec.
         final Intent intent = new Intent(
@@ -876,17 +850,8 @@
         msgView.setText(msg);
         imageView.setImageResource(drawable);
 
-        content.animate().cancel();  // cancel any ongoing animations
-
-        content.setAlpha(0);
         mEmptyView.setVisibility(View.VISIBLE);
         mRecView.setVisibility(View.GONE);
-
-        // fade in the content, so it looks purdy like
-        content.animate()
-                .alpha(1f)
-                .setDuration(EMPTY_REVEAL_DURATION)
-                .setListener(null);
     }
 
     private void showDirectory() {
@@ -1289,12 +1254,14 @@
          * @return The adapter position of the destination item. Could be RecyclerView.NO_POSITION.
          */
         private int findTargetPosition(View view, int keyCode) {
-            if (keyCode == KeyEvent.KEYCODE_MOVE_HOME) {
-                return 0;
-            }
-
-            if (keyCode == KeyEvent.KEYCODE_MOVE_END) {
-                return mAdapter.getItemCount() - 1;
+            switch (keyCode) {
+                case KeyEvent.KEYCODE_MOVE_HOME:
+                    return 0;
+                case KeyEvent.KEYCODE_MOVE_END:
+                    return mAdapter.getItemCount() - 1;
+                case KeyEvent.KEYCODE_PAGE_UP:
+                case KeyEvent.KEYCODE_PAGE_DOWN:
+                    return findTargetPositionByPage(view, keyCode);
             }
 
             // Find a navigation target based on the arrow key that the user pressed.
@@ -1330,6 +1297,50 @@
         }
 
         /**
+         * Given a PgUp/PgDn event and the current view, find the position of the target view.
+         * This returns:
+         * <li>The position of the topmost (or bottom-most) visible item, if the current item is not
+         *     the top- or bottom-most visible item.
+         * <li>The position of an item that is one page's worth of items up (or down) if the current
+         *      item is the top- or bottom-most visible item.
+         * <li>The first (or last) item, if paging up (or down) would go past those limits.
+         * @param view The view that received the key event.
+         * @param keyCode Must be KEYCODE_PAGE_UP or KEYCODE_PAGE_DOWN.
+         * @return The adapter position of the target item.
+         */
+        private int findTargetPositionByPage(View view, int keyCode) {
+            int first = mLayout.findFirstVisibleItemPosition();
+            int last = mLayout.findLastVisibleItemPosition();
+            int current = mRecView.getChildAdapterPosition(view);
+            int pageSize = last - first + 1;
+
+            if (keyCode == KeyEvent.KEYCODE_PAGE_UP) {
+                if (current > first) {
+                    // If the current item isn't the first item, target the first item.
+                    return first;
+                } else {
+                    // If the current item is the first item, target the item one page up.
+                    int target = current - pageSize;
+                    return target < 0 ? 0 : target;
+                }
+            }
+
+            if (keyCode == KeyEvent.KEYCODE_PAGE_DOWN) {
+                if (current < last) {
+                    // If the current item isn't the last item, target the last item.
+                    return last;
+                } else {
+                    // If the current item is the last item, target the item one page down.
+                    int target = current + pageSize;
+                    int max = mAdapter.getItemCount() - 1;
+                    return target < max ? target : max;
+                }
+            }
+
+            throw new IllegalArgumentException("Unsupported keyCode: " + keyCode);
+        }
+
+        /**
          * Requests focus for the item in the given adapter position, scrolling the RecyclerView if
          * necessary.
          *
diff --git a/src/com/android/documentsui/dirlist/DocumentsAdapter.java b/src/com/android/documentsui/dirlist/DocumentsAdapter.java
index 43c2f63..0930c22 100644
--- a/src/com/android/documentsui/dirlist/DocumentsAdapter.java
+++ b/src/com/android/documentsui/dirlist/DocumentsAdapter.java
@@ -27,7 +27,6 @@
 
 import com.android.documentsui.State;
 
-import java.nio.channels.UnsupportedAddressTypeException;
 import java.util.List;
 
 /**
@@ -87,7 +86,7 @@
      * we adjust sizes.
      */
     GridLayoutManager.SpanSizeLookup createSpanSizeLookup() {
-        throw new UnsupportedAddressTypeException();
+        throw new UnsupportedOperationException();
     }
 
     static boolean isDirectory(Cursor cursor) {
diff --git a/tests/src/com/android/documentsui/ActivityTest.java b/tests/src/com/android/documentsui/ActivityTest.java
new file mode 100644
index 0000000..34f1120
--- /dev/null
+++ b/tests/src/com/android/documentsui/ActivityTest.java
@@ -0,0 +1,136 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.documentsui;
+
+import static com.android.documentsui.StubProvider.DEFAULT_AUTHORITY;
+import static com.android.documentsui.StubProvider.ROOT_0_ID;
+import static com.android.documentsui.StubProvider.ROOT_1_ID;
+
+import android.app.Activity;
+import android.app.Instrumentation;
+import android.content.ContentProviderClient;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.content.Intent;
+import android.os.RemoteException;
+import android.provider.DocumentsContract.Document;
+import android.support.test.uiautomator.By;
+import android.support.test.uiautomator.Configurator;
+import android.support.test.uiautomator.UiDevice;
+import android.support.test.uiautomator.UiObjectNotFoundException;
+import android.support.test.uiautomator.Until;
+import android.test.ActivityInstrumentationTestCase2;
+import android.view.MotionEvent;
+
+import com.android.documentsui.model.RootInfo;
+
+/**
+ * Provides basic test environment for UI tests:
+ * - Launches activity
+ * - Creates and gives access to test root directories and test files
+ * - Cleans up the test environment
+ */
+public abstract class ActivityTest<T extends Activity> extends ActivityInstrumentationTestCase2<T> {
+
+    static final int TIMEOUT = 5000;
+
+    // Testing files. For custom ones, override initTestFiles().
+    public static final String dirName1 = "Dir1";
+    public static final String fileName1 = "file1.log";
+    public static final String fileName2 = "file12.png";
+    public static final String fileName3 = "anotherFile0.log";
+    public static final String fileName4 = "poodles.text";
+    public static final String fileNameNoRename = "NO_RENAMEfile.txt";
+
+    public UiBot bot;
+    public UiDevice device;
+    public Context context;
+
+    public RootInfo rootDir0;
+    public RootInfo rootDir1;
+
+    ContentResolver mResolver;
+    DocumentsProviderHelper mDocsHelper;
+    ContentProviderClient mClient;
+
+    public ActivityTest(Class<T> activityClass) {
+        super(activityClass);
+    }
+
+    @Override
+    public void setUp() throws Exception {
+        device = UiDevice.getInstance(getInstrumentation());
+        // NOTE: Must be the "target" context, else security checks in content provider will fail.
+        context = getInstrumentation().getTargetContext();
+        bot = new UiBot(device, context, TIMEOUT);
+
+        Configurator.getInstance().setToolType(MotionEvent.TOOL_TYPE_MOUSE);
+        bot.revealLauncher();
+
+        mResolver = context.getContentResolver();
+        mClient = mResolver.acquireUnstableContentProviderClient(DEFAULT_AUTHORITY);
+        mDocsHelper = new DocumentsProviderHelper(DEFAULT_AUTHORITY, mClient);
+
+        rootDir0 = mDocsHelper.getRoot(ROOT_0_ID);
+        rootDir1 = mDocsHelper.getRoot(ROOT_1_ID);
+
+        launchActivity();
+
+        bot.revealApp();
+        resetStorage();
+    }
+
+    @Override
+    public void tearDown() throws Exception {
+        mClient.release();
+        super.tearDown();
+    }
+
+    void launchActivity() {
+        final Intent intent = context.getPackageManager().getLaunchIntentForPackage(
+                UiBot.TARGET_PKG);
+        intent.addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK | Intent.FLAG_ACTIVITY_NEW_TASK);
+        setActivityIntent(intent);
+        getActivity();  // Launch the activity.
+    }
+
+    void resetStorage() throws RemoteException {
+        mClient.call("clear", null, null);
+        device.waitForIdle();
+    }
+
+    void initTestFiles() throws RemoteException {
+        mDocsHelper.createFolder(rootDir0, dirName1);
+        mDocsHelper.createDocument(rootDir0, "text/plain", fileName1);
+        mDocsHelper.createDocument(rootDir0, "image/png", fileName2);
+        mDocsHelper.createDocumentWithFlags(rootDir0.documentId, "text/plain", fileNameNoRename,
+                Document.FLAG_SUPPORTS_WRITE);
+
+        mDocsHelper.createDocument(rootDir1, "text/plain", fileName3);
+        mDocsHelper.createDocument(rootDir1, "text/plain", fileName4);
+    }
+
+    void assertDefaultContentOfTestDir0() throws UiObjectNotFoundException {
+        bot.assertDocumentsCount(ROOT_0_ID, 4);
+        bot.assertHasDocuments(fileName1, fileName2, dirName1, fileNameNoRename);
+    }
+
+    void assertDefaultContentOfTestDir1() throws UiObjectNotFoundException {
+        bot.assertDocumentsCount(ROOT_1_ID, 2);
+        bot.assertHasDocuments(fileName3, fileName4);
+    }
+}
diff --git a/tests/src/com/android/documentsui/DocumentsProviderHelper.java b/tests/src/com/android/documentsui/DocumentsProviderHelper.java
index 5df4dca..af478ea 100644
--- a/tests/src/com/android/documentsui/DocumentsProviderHelper.java
+++ b/tests/src/com/android/documentsui/DocumentsProviderHelper.java
@@ -92,7 +92,8 @@
             Uri uri = DocumentsContract.createDocument(mClient, parentUri, mimeType, name);
             return uri;
         } catch (RemoteException e) {
-            throw new RuntimeException("Couldn't create document: " + name + " with mimetype " + mimeType, e);
+            throw new RuntimeException("Couldn't create document: " + name + " with mimetype "
+                    + mimeType, e);
         }
     }
 
diff --git a/tests/src/com/android/documentsui/DownloadsActivityUiTest.java b/tests/src/com/android/documentsui/DownloadsActivityUiTest.java
index 243c357..737b376 100644
--- a/tests/src/com/android/documentsui/DownloadsActivityUiTest.java
+++ b/tests/src/com/android/documentsui/DownloadsActivityUiTest.java
@@ -38,122 +38,74 @@
 import com.android.documentsui.model.RootInfo;
 
 @LargeTest
-public class DownloadsActivityUiTest extends ActivityInstrumentationTestCase2<DownloadsActivity> {
+public class DownloadsActivityUiTest extends ActivityTest<DownloadsActivity> {
 
     private static final int TIMEOUT = 5000;
     private static final String TAG = "DownloadsActivityUiTest";
     private static final String TARGET_PKG = "com.android.documentsui";
     private static final String LAUNCHER_PKG = "com.android.launcher";
 
-    private UiBot mBot;
-    private UiDevice mDevice;
-    private Context mContext;
-    private ContentResolver mResolver;
-    private DocumentsProviderHelper mDocsHelper;
-    private ContentProviderClient mClient;
-    private RootInfo mRoot;
-
     public DownloadsActivityUiTest() {
         super(DownloadsActivity.class);
     }
 
-    public void setUp() throws Exception {
-        // Initialize UiDevice instance.
-        Instrumentation instrumentation = getInstrumentation();
-
-        mDevice = UiDevice.getInstance(instrumentation);
-
-        Configurator.getInstance().setToolType(MotionEvent.TOOL_TYPE_MOUSE);
-
-        // Start from the home screen.
-        mDevice.pressHome();
-        mDevice.wait(Until.hasObject(By.pkg(LAUNCHER_PKG).depth(0)), TIMEOUT);
-
-        // NOTE: Must be the "target" context, else security checks in content provider will fail.
-        mContext = instrumentation.getTargetContext();
-        mResolver = mContext.getContentResolver();
-
-        mClient = mResolver.acquireUnstableContentProviderClient(DEFAULT_AUTHORITY);
-        mDocsHelper = new DocumentsProviderHelper(DEFAULT_AUTHORITY, mClient);
-
-        mRoot = mDocsHelper.getRoot(ROOT_0_ID);
-
-        // Open the Downloads activity on our stub provider root.
-        Intent intent = new Intent(DocumentsContract.ACTION_MANAGE_ROOT);
-        intent.setDataAndType(mRoot.getUri(), DocumentsContract.Root.MIME_TYPE_ITEM);
+    @Override
+    void launchActivity() {
+        final Intent intent = new Intent(DocumentsContract.ACTION_MANAGE_ROOT);
+        intent.setDataAndType(rootDir0.getUri(), DocumentsContract.Root.MIME_TYPE_ITEM);
         intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         intent.addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK);
         setActivityIntent(intent);
-        getActivity();  // Start the activity.
-
-        // Wait for the app to appear.
-        mDevice.wait(Until.hasObject(By.pkg(TARGET_PKG).depth(0)), TIMEOUT);
-        mDevice.waitForIdle();
-
-        mBot = new UiBot(mDevice, mContext, TIMEOUT);
-
-        resetStorage();  // Just in case a test failed and tearDown didn't happen.
+        getActivity();  // Launch the activity.
     }
 
     @Override
-    protected void tearDown() throws Exception {
-        Log.d(TAG, "Resetting storage from setUp");
-        resetStorage();
-        mClient.release();
-
-        super.tearDown();
-    }
-
-    private void resetStorage() throws RemoteException {
-        mClient.call("clear", null, null);
-        // TODO: Would be nice to have an event to wait on here.
-        mDevice.waitForIdle();
-    }
-
-    private void initTestFiles() throws RemoteException {
-        mDocsHelper.createDocument(mRoot, "text/plain", "file0.log");
-        mDocsHelper.createDocument(mRoot, "image/png", "file1.png");
-        mDocsHelper.createDocument(mRoot, "text/csv", "file2.csv");
+    void initTestFiles() throws RemoteException {
+        mDocsHelper.createDocument(rootDir0, "text/plain", "file0.log");
+        mDocsHelper.createDocument(rootDir0, "image/png", "file1.png");
+        mDocsHelper.createDocument(rootDir0, "text/csv", "file2.csv");
     }
 
     public void testWindowTitle() throws Exception {
         initTestFiles();
 
-        mBot.assertWindowTitle(ROOT_0_ID);
+        bot.assertWindowTitle(ROOT_0_ID);
     }
 
     public void testFilesListed() throws Exception {
         initTestFiles();
 
-        mBot.assertHasDocuments("file0.log", "file1.png", "file2.csv");
+        bot.assertHasDocuments("file0.log", "file1.png", "file2.csv");
     }
 
     public void testFilesList_LiveUpdate() throws Exception {
         initTestFiles();
 
-        mDocsHelper.createDocument(mRoot, "yummers/sandwich", "Ham & Cheese.sandwich");
-        mBot.assertHasDocuments("file0.log", "file1.png", "file2.csv", "Ham & Cheese.sandwich");
+        mDocsHelper.createDocument(rootDir0, "yummers/sandwich", "Ham & Cheese.sandwich");
+
+        device.waitForIdle();
+        bot.assertHasDocuments("file0.log", "file1.png", "file2.csv", "Ham & Cheese.sandwich");
     }
 
     public void testDeleteDocument() throws Exception {
         initTestFiles();
 
-        mBot.clickDocument("file1.png");
-        mDevice.waitForIdle();
-        mBot.menuDelete().click();
+        bot.clickDocument("file1.png");
+        device.waitForIdle();
+        bot.menuDelete().click();
 
-        mBot.waitForDeleteSnackbar();
-        assertFalse(mBot.hasDocuments("file1.png"));
+        bot.waitForDeleteSnackbar();
+        assertFalse(bot.hasDocuments("file1.png"));
 
-        mBot.waitForDeleteSnackbarGone();
-        assertFalse(mBot.hasDocuments("file1.png"));
+        bot.waitForDeleteSnackbarGone();
+        assertFalse(bot.hasDocuments("file1.png"));
     }
 
     public void testSupportsShare() throws Exception {
         initTestFiles();
 
-        mBot.clickDocument("file1.png");
-        mDevice.waitForIdle();
-        assertNotNull(mBot.menuShare());
+        bot.clickDocument("file1.png");
+        device.waitForIdle();
+        assertNotNull(bot.menuShare());
     }
 }
diff --git a/tests/src/com/android/documentsui/FilesActivityUiTest.java b/tests/src/com/android/documentsui/FilesActivityUiTest.java
index 868dbbb..7fd2416 100644
--- a/tests/src/com/android/documentsui/FilesActivityUiTest.java
+++ b/tests/src/com/android/documentsui/FilesActivityUiTest.java
@@ -30,7 +30,6 @@
 import android.support.test.uiautomator.Configurator;
 import android.support.test.uiautomator.UiDevice;
 import android.support.test.uiautomator.Until;
-import android.test.ActivityInstrumentationTestCase2;
 import android.test.suitebuilder.annotation.LargeTest;
 import android.util.Log;
 import android.view.MotionEvent;
@@ -38,96 +37,35 @@
 import com.android.documentsui.model.RootInfo;
 
 @LargeTest
-public class FilesActivityUiTest extends ActivityInstrumentationTestCase2<FilesActivity> {
+public class FilesActivityUiTest extends ActivityTest<FilesActivity> {
 
     private static final int TIMEOUT = 5000;
     private static final String TAG = "FilesActivityUiTest";
     private static final String TARGET_PKG = "com.android.documentsui";
     private static final String LAUNCHER_PKG = "com.android.launcher";
 
-    private UiBot mBot;
-    private UiDevice mDevice;
-    private Context mContext;
-    private ContentResolver mResolver;
-    private DocumentsProviderHelper mDocsHelper;
-    private ContentProviderClient mClient;
-    private RootInfo mRoot_0;
-    private RootInfo mRoot_1;
-
     public FilesActivityUiTest() {
         super(FilesActivity.class);
     }
 
-    public void setUp() throws Exception {
-        // Initialize UiDevice instance.
-        Instrumentation instrumentation = getInstrumentation();
-
-        mDevice = UiDevice.getInstance(instrumentation);
-
-        Configurator.getInstance().setToolType(MotionEvent.TOOL_TYPE_MOUSE);
-
-        // Start from the home screen.
-        mDevice.pressHome();
-        mDevice.wait(Until.hasObject(By.pkg(LAUNCHER_PKG).depth(0)), TIMEOUT);
-
-        // NOTE: Must be the "target" context, else security checks in content provider will fail.
-        mContext = instrumentation.getTargetContext();
-        mResolver = mContext.getContentResolver();
-
-        mClient = mResolver.acquireUnstableContentProviderClient(DEFAULT_AUTHORITY);
-        assertNotNull("Failed to acquire ContentProviderClient.", mClient);
-        mDocsHelper = new DocumentsProviderHelper(DEFAULT_AUTHORITY, mClient);
-
-        // Launch app.
-        Intent intent = mContext.getPackageManager().getLaunchIntentForPackage(TARGET_PKG);
-        intent.addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK);
-        setActivityIntent(intent);
-        getActivity();  // Start the activity.
-
-        // Wait for the app to appear.
-        mDevice.wait(Until.hasObject(By.pkg(TARGET_PKG).depth(0)), TIMEOUT);
-        mDevice.waitForIdle();
-
-        mBot = new UiBot(mDevice, mContext, TIMEOUT);
-
-        resetStorage();  // Just incase a test failed and tearDown didn't happen.
-    }
-
     @Override
-    protected void tearDown() throws Exception {
-        Log.d(TAG, "Resetting storage from setUp");
-        resetStorage();
-        mClient.release();
+    public void initTestFiles() throws RemoteException {
+        mDocsHelper.createDocument(rootDir0, "text/plain", "file0.log");
+        mDocsHelper.createDocument(rootDir0, "image/png", "file1.png");
+        mDocsHelper.createDocument(rootDir0, "text/csv", "file2.csv");
 
-        super.tearDown();
-    }
-
-    private void resetStorage() throws RemoteException {
-        mClient.call("clear", null, null);
-        // TODO: Would be nice to have an event to wait on here.
-        mDevice.waitForIdle();
-    }
-
-    private void initTestFiles() throws RemoteException {
-        mRoot_0 = mDocsHelper.getRoot(ROOT_0_ID);
-        mRoot_1 = mDocsHelper.getRoot(ROOT_1_ID);
-
-        mDocsHelper.createDocument(mRoot_0, "text/plain", "file0.log");
-        mDocsHelper.createDocument(mRoot_0, "image/png", "file1.png");
-        mDocsHelper.createDocument(mRoot_0, "text/csv", "file2.csv");
-
-        mDocsHelper.createDocument(mRoot_1, "text/plain", "anotherFile0.log");
-        mDocsHelper.createDocument(mRoot_1, "text/plain", "poodles.text");
+        mDocsHelper.createDocument(rootDir1, "text/plain", "anotherFile0.log");
+        mDocsHelper.createDocument(rootDir1, "text/plain", "poodles.text");
     }
 
     public void testRootsListed() throws Exception {
         initTestFiles();
 
-        mBot.openRoot(ROOT_0_ID);
+        bot.openRoot(ROOT_0_ID);
 
         // Should also have Drive, but that requires pre-configuration of devices
         // We omit for now.
-        mBot.assertHasRoots(
+        bot.assertHasRoots(
                 "Images",
                 "Videos",
                 "Audio",
@@ -140,60 +78,60 @@
     public void testFilesListed() throws Exception {
         initTestFiles();
 
-        mBot.openRoot(ROOT_0_ID);
-        mBot.assertHasDocuments("file0.log", "file1.png", "file2.csv");
+        bot.openRoot(ROOT_0_ID);
+        bot.assertHasDocuments("file0.log", "file1.png", "file2.csv");
     }
 
     public void testLoadsHomeByDefault() throws Exception {
         initTestFiles();
 
-        mDevice.waitForIdle();
-        mBot.assertWindowTitle("Home");
+        device.waitForIdle();
+        bot.assertWindowTitle("Home");
     }
 
     public void testRootClickSetsWindowTitle() throws Exception {
         initTestFiles();
 
-        mBot.openRoot("Downloads");
-        mBot.assertWindowTitle("Downloads");
+        bot.openRoot("Downloads");
+        bot.assertWindowTitle("Downloads");
     }
 
     public void testFilesList_LiveUpdate() throws Exception {
         initTestFiles();
 
-        mBot.openRoot(ROOT_0_ID);
-        mDocsHelper.createDocument(mRoot_0, "yummers/sandwich", "Ham & Cheese.sandwich");
+        bot.openRoot(ROOT_0_ID);
+        mDocsHelper.createDocument(rootDir0, "yummers/sandwich", "Ham & Cheese.sandwich");
 
-        mDevice.waitForIdle();
-        mBot.assertHasDocuments("file0.log", "file1.png", "file2.csv", "Ham & Cheese.sandwich");
+        device.waitForIdle();
+        bot.assertHasDocuments("file0.log", "file1.png", "file2.csv", "Ham & Cheese.sandwich");
     }
 
     public void testDeleteDocument() throws Exception {
         initTestFiles();
 
-        mBot.openRoot(ROOT_0_ID);
+        bot.openRoot(ROOT_0_ID);
 
-        mBot.clickDocument("file1.png");
-        mDevice.waitForIdle();
-        mBot.menuDelete().click();
+        bot.clickDocument("file1.png");
+        device.waitForIdle();
+        bot.menuDelete().click();
 
-        mBot.waitForDeleteSnackbar();
-        assertFalse(mBot.hasDocuments("file1.png"));
+        bot.waitForDeleteSnackbar();
+        assertFalse(bot.hasDocuments("file1.png"));
 
-        mBot.waitForDeleteSnackbarGone();
-        assertFalse(mBot.hasDocuments("file1.png"));
+        bot.waitForDeleteSnackbarGone();
+        assertFalse(bot.hasDocuments("file1.png"));
 
         // Now delete from another root.
-        mBot.openRoot(ROOT_1_ID);
+        bot.openRoot(ROOT_1_ID);
 
-        mBot.clickDocument("poodles.text");
-        mDevice.waitForIdle();
-        mBot.menuDelete().click();
+        bot.clickDocument("poodles.text");
+        device.waitForIdle();
+        bot.menuDelete().click();
 
-        mBot.waitForDeleteSnackbar();
-        assertFalse(mBot.hasDocuments("poodles.text"));
+        bot.waitForDeleteSnackbar();
+        assertFalse(bot.hasDocuments("poodles.text"));
 
-        mBot.waitForDeleteSnackbarGone();
-        assertFalse(mBot.hasDocuments("poodles.text"));
+        bot.waitForDeleteSnackbarGone();
+        assertFalse(bot.hasDocuments("poodles.text"));
     }
 }
diff --git a/tests/src/com/android/documentsui/RenameDocumentUiTest.java b/tests/src/com/android/documentsui/RenameDocumentUiTest.java
index 5c6254f..303f2d1 100644
--- a/tests/src/com/android/documentsui/RenameDocumentUiTest.java
+++ b/tests/src/com/android/documentsui/RenameDocumentUiTest.java
@@ -17,7 +17,6 @@
 package com.android.documentsui;
 
 import static com.android.documentsui.StubProvider.ROOT_0_ID;
-import static com.android.documentsui.UiTestEnvironment.TIMEOUT;
 
 import android.support.test.uiautomator.UiObject;
 import android.support.test.uiautomator.UiObjectNotFoundException;
@@ -25,142 +24,141 @@
 import android.test.suitebuilder.annotation.LargeTest;
 
 @LargeTest
-public class RenameDocumentUiTest extends InstrumentationTestCase {
+public class RenameDocumentUiTest extends ActivityTest<FilesActivity> {
 
     private static final String TAG = "RenamDocumentUiTest";
 
     private final String newName = "kitties.log";
 
-    private UiTestEnvironment mHelper;
+    public RenameDocumentUiTest() {
+        super(FilesActivity.class);
+    }
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
-        mHelper = new UiTestEnvironment(getInstrumentation());
-        mHelper.launch();
-        mHelper.initTestFiles();
-        mHelper.bot().openRoot(ROOT_0_ID);
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
     }
 
     public void testRenameEnabled_SingleSelection() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().assertMenuEnabled(R.string.menu_rename, true);
+        bot.selectDocument(fileName1);
+        bot.openOverflowMenu();
+        bot.assertMenuEnabled(R.string.menu_rename, true);
 
         // Dismiss more options window
-        mHelper.device().pressBack();
+        device.pressBack();
     }
 
     public void testNoRenameSupport_SingleSelection() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileNameNoRename);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().assertMenuEnabled(R.string.menu_rename, false);
+        bot.selectDocument(fileNameNoRename);
+        bot.openOverflowMenu();
+        bot.assertMenuEnabled(R.string.menu_rename, false);
 
         // Dismiss more options window
-        mHelper.device().pressBack();
+        device.pressBack();
     }
 
     public void testOneHasRenameSupport_MultipleSelection() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().selectDocument(UiTestEnvironment.fileNameNoRename);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().assertMenuEnabled(R.string.menu_rename, false);
+        bot.selectDocument(fileName1);
+        bot.selectDocument(fileNameNoRename);
+        bot.openOverflowMenu();
+        bot.assertMenuEnabled(R.string.menu_rename, false);
 
         // Dismiss more options window
-        mHelper.device().pressBack();
+        device.pressBack();
     }
 
     public void testRenameDisabled_MultipleSelection() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName2);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().assertMenuEnabled(R.string.menu_rename, false);
+        bot.selectDocument(fileName1);
+        bot.selectDocument(fileName2);
+        bot.openOverflowMenu();
+        bot.assertMenuEnabled(R.string.menu_rename, false);
 
         // Dismiss more options window
-        mHelper.device().pressBack();
+        device.pressBack();
     }
 
     public void testRenameFile_OkButton() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().openDialog(R.string.menu_rename);
-        mHelper.bot().setDialogText(newName);
-        mHelper.bot().dismissKeyboardIfPresent();
+        bot.selectDocument(fileName1);
+        bot.openOverflowMenu();
+        bot.openDialog(R.string.menu_rename);
+        bot.setDialogText(newName);
+        bot.dismissKeyboardIfPresent();
 
-        mHelper.device().waitForIdle(TIMEOUT);
-        mHelper.bot().findRenameDialogOkButton().click();
-        mHelper.device().waitForIdle(TIMEOUT);
+        device.waitForIdle(TIMEOUT);
+        bot.findRenameDialogOkButton().click();
+        device.waitForIdle(TIMEOUT);
 
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName1, false);
-        mHelper.bot().assertDocument(newName, true);
-        mHelper.bot().assertDocumentsCount(mHelper.getDocumentsCountDir0());
+        bot.assertDocument(fileName1, false);
+        bot.assertDocument(newName, true);
+        bot.assertDocumentsCount(4);
     }
 
     public void testRenameFile_Enter() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().openDialog(R.string.menu_rename);
-        mHelper.bot().setDialogText(newName);
+        bot.selectDocument(fileName1);
+        bot.openOverflowMenu();
+        bot.openDialog(R.string.menu_rename);
+        bot.setDialogText(newName);
 
         pressEnter();
 
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName1, false);
-        mHelper.bot().assertDocument(newName, true);
-        mHelper.bot().assertDocumentsCount(mHelper.getDocumentsCountDir0());
+        bot.assertDocument(fileName1, false);
+        bot.assertDocument(newName, true);
+        bot.assertDocumentsCount(4);
     }
 
     public void testRenameFile_Cancel() throws Exception {
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().openDialog(R.string.menu_rename);
-        mHelper.bot().setDialogText(newName);
-        mHelper.bot().dismissKeyboardIfPresent();
+        bot.selectDocument(fileName1);
+        bot.openOverflowMenu();
+        bot.openDialog(R.string.menu_rename);
+        bot.setDialogText(newName);
+        bot.dismissKeyboardIfPresent();
 
-        mHelper.device().waitForIdle(TIMEOUT);
-        mHelper.bot().findRenameDialogCancelButton().click();
-        mHelper.device().waitForIdle(TIMEOUT);
+        device.waitForIdle(TIMEOUT);
+        bot.findRenameDialogCancelButton().click();
+        device.waitForIdle(TIMEOUT);
 
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName1, true);
-        mHelper.bot().assertDocument(newName, false);
-        mHelper.bot().assertDocumentsCount(mHelper.getDocumentsCountDir0());
+        bot.assertDocument(fileName1, true);
+        bot.assertDocument(newName, false);
+        bot.assertDocumentsCount(4);
     }
 
     public void testRenameDir() throws Exception {
         String oldName = "Dir1";
         String newName = "Dir123";
 
-        mHelper.bot().selectDocument(oldName);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().openDialog(R.string.menu_rename);
-        mHelper.bot().setDialogText(newName);
+        bot.selectDocument(oldName);
+        bot.openOverflowMenu();
+        bot.openDialog(R.string.menu_rename);
+        bot.setDialogText(newName);
 
         pressEnter();
 
-        mHelper.bot().assertDocument(oldName, false);
-        mHelper.bot().assertDocument(newName, true);
-        mHelper.bot().assertDocumentsCount(mHelper.getDocumentsCountDir0());
+        bot.assertDocument(oldName, false);
+        bot.assertDocument(newName, true);
+        bot.assertDocumentsCount(4);
     }
 
     public void testRename_NameExists() throws Exception {
         // Check that document with the new name exists
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName2, true);
-        mHelper.bot().selectDocument(UiTestEnvironment.fileName1);
-        mHelper.bot().openOverflowMenu();
-        mHelper.bot().openDialog(R.string.menu_rename);
-        mHelper.bot().setDialogText(UiTestEnvironment.fileName2);
+        bot.assertDocument(fileName2, true);
+        bot.selectDocument(fileName1);
+        bot.openOverflowMenu();
+        bot.openDialog(R.string.menu_rename);
+        bot.setDialogText(fileName2);
 
         pressEnter();
 
-        mHelper.bot().assertSnackbar(R.string.rename_error);
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName1, true);
-        mHelper.bot().assertDocument(UiTestEnvironment.fileName2, true);
-        mHelper.bot().assertDocumentsCount(mHelper.getDocumentsCountDir0());
+        bot.assertSnackbar(R.string.rename_error);
+        bot.assertDocument(fileName1, true);
+        bot.assertDocument(fileName2, true);
+        bot.assertDocumentsCount(4);
     }
 
     private void pressEnter() {
-        mHelper.device().waitForIdle(TIMEOUT);
-        mHelper.device().pressEnter();
-        mHelper.device().waitForIdle(TIMEOUT);
+        device.waitForIdle(TIMEOUT);
+        device.pressEnter();
+        device.waitForIdle(TIMEOUT);
     }
-
 }
diff --git a/tests/src/com/android/documentsui/SearchViewUiTest.java b/tests/src/com/android/documentsui/SearchViewUiTest.java
index 042ec85..61da3df 100644
--- a/tests/src/com/android/documentsui/SearchViewUiTest.java
+++ b/tests/src/com/android/documentsui/SearchViewUiTest.java
@@ -24,138 +24,132 @@
 import android.test.suitebuilder.annotation.LargeTest;
 
 @LargeTest
-public class SearchViewUiTest extends InstrumentationTestCase {
+public class SearchViewUiTest extends ActivityTest<FilesActivity> {
 
     private static final String TAG = "SearchViewUiTest";
 
-    private UiTestEnvironment mEnv;
-
-    @Override
-    public void setUp() throws Exception {
-        super.setUp();
-        mEnv = new UiTestEnvironment(getInstrumentation());
-        mEnv.launch();
-
-    }
-
-    @Override
-    protected void tearDown() throws Exception {
-        mEnv.device().pressBack();
-        super.tearDown();
+    public SearchViewUiTest() {
+        super(FilesActivity.class);
     }
 
     public void testSearchView_ExpandsOnClick() throws Exception {
-        mEnv.bot().openSearchView();
-        mEnv.bot().assertSearchTextFiledAndIcon(true, false);
+        bot.openSearchView();
+        bot.assertSearchTextFiledAndIcon(true, false);
     }
 
     public void testSearchView_CollapsesOnBack() throws Exception {
-        mEnv.bot().openSearchView();
+        bot.openSearchView();
 
-        mEnv.device().pressBack();
+        device.pressBack();
 
-        mEnv.bot().assertSearchTextFiledAndIcon(false, true);
+        bot.assertSearchTextFiledAndIcon(false, true);
     }
 
     public void testSearchView_ClearsTextOnBack() throws Exception {
         String query = "file2";
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
+        bot.openSearchView();
+        bot.setSearchQuery(query);
 
-        mEnv.device().pressBack();
+        device.pressBack();
 
-        mEnv.bot().assertSearchTextFiledAndIcon(false, true);
+        bot.assertSearchTextFiledAndIcon(false, true);
     }
 
     public void testSearch_ResultsFound() throws Exception {
-        mEnv.initTestFiles();
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
 
         String query = "file1";
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
-        mEnv.bot().assertSearchTextField(true, query);
+        bot.openSearchView();
+        bot.setSearchQuery(query);
+        bot.assertSearchTextField(true, query);
 
-        mEnv.device().pressEnter();
+        device.pressEnter();
 
-        mEnv.bot().assertDocumentsCountOnList(true, 2);
-        mEnv.bot().assertHasDocuments(UiTestEnvironment.fileName1, UiTestEnvironment.fileName2);
-        mEnv.bot().assertSearchTextField(false, query);
+        bot.assertDocumentsCountOnList(true, 2);
+        bot.assertHasDocuments(fileName1, fileName2);
+
+        bot.assertSearchTextField(false, query);
     }
 
     public void testSearchResultsFound_ClearsOnBack() throws Exception {
-        mEnv.initTestFiles();
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
 
-        String query = UiTestEnvironment.fileName1;
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
+        String query = fileName1;
+        bot.openSearchView();
+        bot.setSearchQuery(query);
 
-        mEnv.device().pressEnter();
-        mEnv.device().pressBack();
-        mEnv.assertDefaultContentOfTestDir0();
+        device.pressEnter();
+        device.pressBack();
+
+        assertDefaultContentOfTestDir0();
     }
 
     public void testSearch_NoResults() throws Exception {
-        mEnv.initTestFiles();
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
 
         String query = "chocolate";
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
+        bot.openSearchView();
+        bot.setSearchQuery(query);
 
-        mEnv.device().pressEnter();
+        device.pressEnter();
 
-        mEnv.bot().assertDocumentsCountOnList(false, 0);
+        bot.assertDocumentsCountOnList(false, 0);
 
-        String msg = String.valueOf(mEnv.context().getString(R.string.no_results));
-        mEnv.bot().assertMessageTextView(String.format(msg, "TEST_ROOT_0"));
-        mEnv.bot().assertSearchTextField(false, query);
+        device.waitForIdle();
+        String msg = String.valueOf(context.getString(R.string.no_results));
+        bot.assertMessageTextView(String.format(msg, "TEST_ROOT_0"));
+
+        bot.assertSearchTextField(false, query);
     }
 
     public void testSearchNoResults_ClearsOnBack() throws Exception {
-        mEnv.initTestFiles();
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
 
         String query = "chocolate";
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
+        bot.openSearchView();
+        bot.setSearchQuery(query);
 
-        mEnv.device().pressEnter();
-        mEnv.device().pressBack();
-        mEnv.assertDefaultContentOfTestDir0();
+        device.pressEnter();
+        device.pressBack();
+
+        device.waitForIdle();
+        assertDefaultContentOfTestDir0();
     }
 
     public void testSearchResultsFound_ClearsOnDirectoryChange() throws Exception {
-        mEnv.initTestFiles();
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        initTestFiles();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
 
-        String query = UiTestEnvironment.fileName1;
-        mEnv.bot().openSearchView();
-        mEnv.bot().setSearchQuery(query);
+        String query = fileName1;
+        bot.openSearchView();
+        bot.setSearchQuery(query);
 
-        mEnv.device().pressEnter();
+        device.pressEnter();
 
-        mEnv.bot().openRoot(ROOT_1_ID);
-        mEnv.assertDefaultContentOfTestDir1();
+        bot.openRoot(ROOT_1_ID);
+        assertDefaultContentOfTestDir1();
 
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.assertDefaultContentOfTestDir0();
+        bot.openRoot(ROOT_0_ID);
+        assertDefaultContentOfTestDir0();
     }
 
     public void testSearchIconVisible_RootWithSearchSupport() throws Exception {
-        mEnv.bot().openRoot(ROOT_0_ID);
-        mEnv.bot().assertSearchTextFiledAndIcon(false, true);
+        bot.openRoot(ROOT_0_ID);
+        bot.assertSearchTextFiledAndIcon(false, true);
     }
 
     public void testSearchIconHidden_RootNoSearchSupport() throws Exception {
-        mEnv.bot().openRoot(ROOT_1_ID);
-        mEnv.bot().assertSearchTextFiledAndIcon(false, false);
+        bot.openRoot(ROOT_1_ID);
+        bot.assertSearchTextFiledAndIcon(false, false);
     }
 
 }
diff --git a/tests/src/com/android/documentsui/UiBot.java b/tests/src/com/android/documentsui/UiBot.java
index 4f6300a..6ba2146 100644
--- a/tests/src/com/android/documentsui/UiBot.java
+++ b/tests/src/com/android/documentsui/UiBot.java
@@ -49,8 +49,10 @@
  * programmatically, and making assertions against the state of the UI.
  */
 class UiBot {
+    public static final String TARGET_PKG = "com.android.documentsui";
 
     private static final String TAG = "UiBot";
+    private static final String LAUNCHER_PKG = "com.android.launcher";
 
     private static final BySelector SNACK_DELETE =
             By.desc(Pattern.compile("^Deleting [0-9]+ file.+"));
@@ -400,4 +402,13 @@
         }
     }
 
+    void revealLauncher() {
+        mDevice.pressHome();
+        mDevice.wait(Until.hasObject(By.pkg(LAUNCHER_PKG).depth(0)), mTimeout);
+    }
+
+    void revealApp() {
+        mDevice.wait(Until.hasObject(By.pkg(TARGET_PKG).depth(0)), mTimeout);
+        mDevice.waitForIdle();
+    }
 }
diff --git a/tests/src/com/android/documentsui/UiTestEnvironment.java b/tests/src/com/android/documentsui/UiTestEnvironment.java
deleted file mode 100644
index 9e30589..0000000
--- a/tests/src/com/android/documentsui/UiTestEnvironment.java
+++ /dev/null
@@ -1,177 +0,0 @@
-/*
- * Copyright (C) 2016 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.documentsui;
-
-import static com.android.documentsui.StubProvider.DEFAULT_AUTHORITY;
-import static com.android.documentsui.StubProvider.ROOT_0_ID;
-import static com.android.documentsui.StubProvider.ROOT_1_ID;
-
-import android.app.Instrumentation;
-import android.content.ContentProviderClient;
-import android.content.ContentResolver;
-import android.content.Context;
-import android.content.Intent;
-import android.os.RemoteException;
-import android.provider.DocumentsContract.Document;
-import android.support.test.uiautomator.By;
-import android.support.test.uiautomator.Configurator;
-import android.support.test.uiautomator.UiDevice;
-import android.support.test.uiautomator.UiObjectNotFoundException;
-import android.support.test.uiautomator.Until;
-import android.view.MotionEvent;
-
-import com.android.documentsui.model.RootInfo;
-
-/**
- * Provides basic test environment for UI tests:
- * - Launches activity
- * - Creates and gives access to test root directories and test files
- * - Cleans up the test environment
- */
-class UiTestEnvironment {
-
-    public static final int TIMEOUT = 5000;
-    public static final String NO_RENAME = "NO_RENAME";
-
-    public static final String dirName1 = "Dir1";
-    public static final String fileName1 = "file1.log";
-    public static final String fileName2 = "file12.png";
-    public static final String fileName3 = "anotherFile0.log";
-    public static final String fileName4 = "poodles.text";
-    public static final String fileNameNoRename = NO_RENAME + "file.txt";
-
-    private static final String TARGET_PKG = "com.android.documentsui";
-    private static final String LAUNCHER_PKG = "com.android.launcher";
-
-    private final UiBot mBot;
-    private final UiDevice mDevice;
-    private final Context mContext;
-
-    private  RootInfo mRootDir0;
-    private  RootInfo mRootDir1;
-    private int mDocsCountDir0;
-    private int mDocsCountDir1;
-
-    private ContentResolver mResolver;
-    private DocumentsProviderHelper mDocsHelper;
-    private ContentProviderClient mClient;
-
-    private final Instrumentation mInstrumentation;
-
-    public UiTestEnvironment(Instrumentation instrumentation) {
-        mInstrumentation = instrumentation;
-        mDevice = UiDevice.getInstance(mInstrumentation);
-        // NOTE: Must be the "target" context, else security checks in content provider will fail.
-        mContext = mInstrumentation.getTargetContext();
-        mBot = new UiBot(mDevice, mContext, TIMEOUT);
-    }
-
-/**
- * Launches default activity and waits for the application to appear.
- * @throws Exception
- */
-    public void launch() throws Exception {
-        Intent intent = mContext.getPackageManager().getLaunchIntentForPackage(TARGET_PKG);
-        launch(intent);
-    }
-
-    /**
-     * Launches activity specified by intent and waits for the application to appear.
-     * @param intent Intent describing activity to launch.
-     * @throws Exception
-     */
-    public void launch(Intent intent) throws Exception {
-        Configurator.getInstance().setToolType(MotionEvent.TOOL_TYPE_MOUSE);
-        // Start from the home screen.
-        mDevice.pressHome();
-        mDevice.wait(Until.hasObject(By.pkg(LAUNCHER_PKG).depth(0)), TIMEOUT);
-
-        mResolver = mContext.getContentResolver();
-        mClient = mResolver.acquireUnstableContentProviderClient(DEFAULT_AUTHORITY);
-        mDocsHelper = new DocumentsProviderHelper(DEFAULT_AUTHORITY, mClient);
-
-        // Launch app.
-        intent.addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK | Intent.FLAG_ACTIVITY_NEW_TASK);
-        mContext.startActivity(intent);
-        // Wait for the app to appear.
-        mDevice.wait(Until.hasObject(By.pkg(TARGET_PKG).depth(0)), TIMEOUT);
-        mDevice.waitForIdle();
-
-        resetStorage(); // Just incase a test failed and tearDown didn't happen.
-    }
-
-    public void cleanUp() throws Exception {
-        resetStorage();
-        mClient.release();
-    }
-
-    public void resetStorage() throws RemoteException {
-        mClient.call("clear", null, null);
-        mDevice.waitForIdle();
-    }
-
-    public void initTestFiles() throws RemoteException {
-        mRootDir0 = mDocsHelper.getRoot(ROOT_0_ID);
-        mRootDir1 = mDocsHelper.getRoot(ROOT_1_ID);
-
-        mDocsHelper.createFolder(mRootDir0, dirName1);
-        mDocsHelper.createDocument(mRootDir0, "text/plain", fileName1);
-        mDocsHelper.createDocument(mRootDir0, "image/png", fileName2);
-        mDocsHelper.createDocumentWithFlags(mRootDir0.documentId, "text/plain", fileNameNoRename,
-                Document.FLAG_SUPPORTS_WRITE);
-        mDocsCountDir0 = 4;
-
-        mDocsHelper.createDocument(mRootDir1, "text/plain", fileName3);
-        mDocsHelper.createDocument(mRootDir1, "text/plain", fileName4);
-        mDocsCountDir1 = 2;
-    }
-
-    public void assertDefaultContentOfTestDir0() throws UiObjectNotFoundException {
-        bot().assertDocumentsCount(ROOT_0_ID, getDocumentsCountDir0());
-        bot().assertHasDocuments(UiTestEnvironment.fileName1, UiTestEnvironment.fileName2,
-                UiTestEnvironment.dirName1, UiTestEnvironment.fileNameNoRename);
-    }
-
-    public void assertDefaultContentOfTestDir1() throws UiObjectNotFoundException {
-        bot().assertDocumentsCount(ROOT_1_ID, getDocumentsCountDir1());
-        bot().assertHasDocuments(UiTestEnvironment.fileName3, UiTestEnvironment.fileName4);
-    }
-
-    public UiBot bot() {
-        return mBot;
-    }
-
-    public UiDevice device() {
-        return mDevice;
-    }
-
-    public Context context() {
-        return mContext;
-    }
-
-    public RootInfo getRootDir0() {
-        return mRootDir0;
-    }
-
-    public int getDocumentsCountDir0() {
-        return mDocsCountDir0;
-    }
-
-    public int getDocumentsCountDir1() {
-        return mDocsCountDir1;
-    }
-}