Add unit tests for picker.ActionHandler.

Bug: 36094136
Change-Id: Ia3fe43ab1046a5cde87732a840fc6c062c8c8d0f

11 files changed, 462 insertions, 84 deletions

diff --git a/src/com/android/documentsui/DocumentsAccess.java b/src/com/android/documentsui/DocumentsAccess.java
index 5d09aa610..90c344a72 100644
--- a/src/com/android/documentsui/DocumentsAccess.java
+++ b/src/com/android/documentsui/DocumentsAccess.java
@@ -51,6 +51,8 @@ public interface DocumentsAccess {
 
     List<DocumentInfo> getDocuments(String authority, List<String> docIds) throws RemoteException;
 
+    @Nullable Uri createDocument(DocumentInfo parentDoc, String mimeType, String displayName);
+
     public static DocumentsAccess create(Context context) {
         return new RuntimeDocumentAccess(context);
     }
@@ -125,5 +127,18 @@ public interface DocumentsAccess {
                 return DocumentsContract.findDocumentPath(client, docUri);
             }
         }
+
+        @Override
+        public Uri createDocument(DocumentInfo parentDoc, String mimeType, String displayName) {
+            final ContentResolver resolver = mContext.getContentResolver();
+            try (ContentProviderClient client = DocumentsApplication.acquireUnstableProviderOrThrow(
+                        resolver, parentDoc.derivedUri.getAuthority())) {
+                return DocumentsContract.createDocument(
+                        client, parentDoc.derivedUri, mimeType, displayName);
+            } catch (Exception e) {
+                Log.w(TAG, "Failed to create document", e);
+                return null;
+            }
+        }
     }
 }
diff --git a/src/com/android/documentsui/base/DocumentStack.java b/src/com/android/documentsui/base/DocumentStack.java
index 114f54570..6a7c6ca86 100644
--- a/src/com/android/documentsui/base/DocumentStack.java
+++ b/src/com/android/documentsui/base/DocumentStack.java
@@ -36,6 +36,7 @@ import java.net.ProtocolException;
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Objects;
 
 import javax.annotation.Nullable;
 
@@ -244,6 +245,26 @@ public class DocumentStack implements Durable, Parcelable {
     }
 
     @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+
+        if (!(o instanceof DocumentStack)) {
+            return false;
+        }
+
+        DocumentStack other = (DocumentStack) o;
+        return Objects.equals(mRoot, other.mRoot)
+                && mList.equals(other.mList);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mRoot, mList);
+    }
+
+    @Override
     public void read(DataInputStream in) throws IOException {
         final int version = in.readInt();
         switch (version) {
diff --git a/src/com/android/documentsui/picker/ActionHandler.java b/src/com/android/documentsui/picker/ActionHandler.java
index fd2243eb0..d6a27291e 100644
--- a/src/com/android/documentsui/picker/ActionHandler.java
+++ b/src/com/android/documentsui/picker/ActionHandler.java
@@ -19,6 +19,7 @@ package com.android.documentsui.picker;
 import static com.android.documentsui.base.Shared.DEBUG;
 import static com.android.documentsui.base.State.ACTION_CREATE;
 import static com.android.documentsui.base.State.ACTION_GET_CONTENT;
+import static com.android.documentsui.base.State.ACTION_OPEN;
 import static com.android.documentsui.base.State.ACTION_OPEN_TREE;
 import static com.android.documentsui.base.State.ACTION_PICK_COPY_DESTINATION;
 
@@ -155,8 +156,7 @@ class ActionHandler<T extends Activity & Addons> extends AbstractActionHandler<T
                 .execute();
     }
 
-    @VisibleForTesting
-    void onLastAccessedStackLoaded(@Nullable DocumentStack stack) {
+    private void onLastAccessedStackLoaded(@Nullable DocumentStack stack) {
         if (stack == null) {
             loadDefaultLocation();
         } else {
@@ -167,12 +167,11 @@ class ActionHandler<T extends Activity & Addons> extends AbstractActionHandler<T
 
     private void loadDefaultLocation() {
         switch (mState.action) {
-            case ACTION_PICK_COPY_DESTINATION:
-            case State.ACTION_CREATE:
+            case ACTION_CREATE:
                 loadHomeDir();
                 break;
             case ACTION_GET_CONTENT:
-            case State.ACTION_OPEN:
+            case ACTION_OPEN:
             case ACTION_OPEN_TREE:
                 mState.stack.changeRoot(mRoots.getRecentsRoot());
                 mActivity.refreshCurrentRootAndDirectory(AnimationView.ANIM_NONE);
@@ -259,6 +258,7 @@ class ActionHandler<T extends Activity & Addons> extends AbstractActionHandler<T
         assert(mState.action == ACTION_CREATE);
         new CreatePickedDocumentTask(
                 mActivity,
+                mDocs,
                 mLastAccessed,
                 mState.stack,
                 mimeType,
@@ -326,12 +326,12 @@ class ActionHandler<T extends Activity & Addons> extends AbstractActionHandler<T
                     | Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION);
         }
 
-        mActivity.setResult(Activity.RESULT_OK, intent);
+        mActivity.setResult(Activity.RESULT_OK, intent, 0);
         mActivity.finish();
     }
 
     private Executor getExecutorForCurrentDirectory() {
-        final DocumentInfo cwd = mActivity.getCurrentDirectory();
+        final DocumentInfo cwd = mState.stack.peek();
         if (cwd != null && cwd.authority != null) {
             return mExecutors.lookup(cwd.authority);
         } else {
@@ -342,5 +342,12 @@ class ActionHandler<T extends Activity & Addons> extends AbstractActionHandler<T
     public interface Addons extends CommonAddons {
         void onAppPicked(ResolveInfo info);
         void onDocumentPicked(DocumentInfo doc);
+
+        /**
+         * Overload final method {@link Activity#setResult(int, Intent)} so that we can intercept
+         * this method call in test environment.
+         */
+        @VisibleForTesting
+        void setResult(int resultCode, Intent result, int notUsed);
     }
 }
diff --git a/src/com/android/documentsui/picker/CreatePickedDocumentTask.java b/src/com/android/documentsui/picker/CreatePickedDocumentTask.java
index 4a96c872d..25c57f27f 100644
--- a/src/com/android/documentsui/picker/CreatePickedDocumentTask.java
+++ b/src/com/android/documentsui/picker/CreatePickedDocumentTask.java
@@ -24,6 +24,7 @@ import android.provider.DocumentsContract;
 import android.support.design.widget.Snackbar;
 import android.util.Log;
 
+import com.android.documentsui.DocumentsAccess;
 import com.android.documentsui.DocumentsApplication;
 import com.android.documentsui.R;
 import com.android.documentsui.base.BooleanConsumer;
@@ -38,9 +39,8 @@ import java.util.function.Consumer;
  * Task that creates a new document in the background.
  */
 class CreatePickedDocumentTask extends PairedTask<Activity, Void, Uri> {
-    private static final String TAG = "CreatePickedDocumentTas";
-
     private final LastAccessedStorage mLastAccessed;
+    private final DocumentsAccess mDocs;
     private final DocumentStack mStack;
     private final String mMimeType;
     private final String mDisplayName;
@@ -49,6 +49,7 @@ class CreatePickedDocumentTask extends PairedTask<Activity, Void, Uri> {
 
     CreatePickedDocumentTask(
             Activity activity,
+            DocumentsAccess docs,
             LastAccessedStorage lastAccessed,
             DocumentStack stack,
             String mimeType,
@@ -57,6 +58,7 @@ class CreatePickedDocumentTask extends PairedTask<Activity, Void, Uri> {
             Consumer<Uri> callback) {
         super(activity);
         mLastAccessed = lastAccessed;
+        mDocs = docs;
         mStack = stack;
         mMimeType = mimeType;
         mDisplayName = displayName;
@@ -73,19 +75,7 @@ class CreatePickedDocumentTask extends PairedTask<Activity, Void, Uri> {
     protected Uri run(Void... params) {
         DocumentInfo cwd = mStack.peek();
 
-        final ContentResolver resolver = mOwner.getContentResolver();
-        ContentProviderClient client = null;
-        Uri childUri = null;
-        try {
-            client = DocumentsApplication.acquireUnstableProviderOrThrow(
-                    resolver, cwd.derivedUri.getAuthority());
-            childUri = DocumentsContract.createDocument(
-                    client, cwd.derivedUri, mMimeType, mDisplayName);
-        } catch (Exception e) {
-            Log.w(TAG, "Failed to create document", e);
-        } finally {
-            ContentProviderClient.releaseQuietly(client);
-        }
+        Uri childUri = mDocs.createDocument(cwd, mMimeType, mDisplayName);
 
         if (childUri != null) {
             mLastAccessed.setLastAccessed(mOwner, mStack);
diff --git a/src/com/android/documentsui/picker/PickActivity.java b/src/com/android/documentsui/picker/PickActivity.java
index e3d245e53..c61f8a92e 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -350,6 +350,11 @@ public class PickActivity extends BaseActivity implements ActionHandler.Addons {
                 || super.onKeyDown(keyCode, event);
     }
 
+    @Override
+    public void setResult(int resultCode, Intent intent, int notUsed) {
+        setResult(resultCode, intent);
+    }
+
     public static PickActivity get(Fragment fragment) {
         return (PickActivity) fragment.getActivity();
     }
diff --git a/tests/common/com/android/documentsui/TestActivity.java b/tests/common/com/android/documentsui/TestActivity.java
index 93e8b4481..73d6495e9 100644
--- a/tests/common/com/android/documentsui/TestActivity.java
+++ b/tests/common/com/android/documentsui/TestActivity.java
@@ -33,6 +33,7 @@ import com.android.documentsui.AbstractActionHandler.CommonAddons;
 import com.android.documentsui.base.DocumentInfo;
 import com.android.documentsui.base.RootInfo;
 import com.android.documentsui.testing.TestEnv;
+import com.android.documentsui.testing.TestEventHandler;
 import com.android.documentsui.testing.TestEventListener;
 import com.android.documentsui.testing.TestLoaderManager;
 import com.android.documentsui.testing.TestPackageManager;
@@ -59,6 +60,7 @@ public abstract class TestActivity extends AbstractBase {
     public TestEventListener<Integer> refreshCurrentRootAndDirectory;
     public TestEventListener<Boolean> setRootsDrawerOpen;
     public TestEventListener<Uri> notifyDirectoryNavigated;
+    public TestEventHandler<Void> finishedHandler;
 
     public static TestActivity create(TestEnv env) {
         TestActivity activity = Mockito.mock(TestActivity.class, Mockito.CALLS_REAL_METHODS);
@@ -66,20 +68,21 @@ public abstract class TestActivity extends AbstractBase {
         return activity;
     }
 
-   public void init(TestEnv env) {
-       resources = TestResources.create();
-       packageMgr = TestPackageManager.create();
-       intent = new Intent();
+    public void init(TestEnv env) {
+        resources = TestResources.create();
+        packageMgr = TestPackageManager.create();
+        intent = new Intent();
 
-       startActivity = new TestEventListener<>();
-       startService = new TestEventListener<>();
-       rootPicked = new TestEventListener<>();
-       refreshCurrentRootAndDirectory =  new TestEventListener<>();
-       setRootsDrawerOpen = new TestEventListener<>();
-       notifyDirectoryNavigated = new TestEventListener<>();
-       contentResolver = env.contentResolver;
-       loaderManager = new TestLoaderManager();
-   }
+        startActivity = new TestEventListener<>();
+        startService = new TestEventListener<>();
+        rootPicked = new TestEventListener<>();
+        refreshCurrentRootAndDirectory =  new TestEventListener<>();
+        setRootsDrawerOpen = new TestEventListener<>();
+        notifyDirectoryNavigated = new TestEventListener<>();
+        contentResolver = env.contentResolver;
+        loaderManager = new TestLoaderManager();
+        finishedHandler = new TestEventHandler<>();
+    }
 
     @Override
     public final String getPackageName() {
@@ -172,6 +175,11 @@ public abstract class TestActivity extends AbstractBase {
     public final LoaderManager getLoaderManager() {
         return loaderManager;
     }
+
+    @Override
+    public final void finish() {
+        finishedHandler.accept(null);
+    }
 }
 
 // Trick Mockito into finding our Addons methods correctly. W/o this
diff --git a/tests/common/com/android/documentsui/testing/TestDocumentsAccess.java b/tests/common/com/android/documentsui/testing/TestDocumentsAccess.java
index 890af2421..266a06b69 100644
--- a/tests/common/com/android/documentsui/testing/TestDocumentsAccess.java
+++ b/tests/common/com/android/documentsui/testing/TestDocumentsAccess.java
@@ -15,9 +15,13 @@
  */
 package com.android.documentsui.testing;
 
+import static junit.framework.Assert.assertEquals;
+
 import android.net.Uri;
 import android.os.RemoteException;
+import android.provider.DocumentsContract;
 import android.provider.DocumentsContract.Path;
+import android.util.Pair;
 
 import com.android.documentsui.DocumentsAccess;
 import com.android.documentsui.base.DocumentInfo;
@@ -38,6 +42,8 @@ public class TestDocumentsAccess implements DocumentsAccess {
 
     public TestEventHandler<Uri> lastUri = new TestEventHandler<>();
 
+    private Pair<DocumentInfo, DocumentInfo> mLastCreatedDoc;
+
     @Override
     public DocumentInfo getRootDocument(RootInfo root) {
         return nextRootDocument;
@@ -54,6 +60,20 @@ public class TestDocumentsAccess implements DocumentsAccess {
     }
 
     @Override
+    public Uri createDocument(DocumentInfo parentDoc, String mimeType, String displayName) {
+        final DocumentInfo child = new DocumentInfo();
+        child.authority = parentDoc.authority;
+        child.mimeType = mimeType;
+        child.displayName = displayName;
+        child.documentId = displayName;
+        child.derivedUri = DocumentsContract.buildDocumentUri(child.authority, displayName);
+
+        mLastCreatedDoc = Pair.create(parentDoc, child);
+
+        return child.derivedUri;
+    }
+
+    @Override
     public DocumentInfo getArchiveDocument(Uri uri) {
         return nextDocument;
     }
@@ -68,4 +88,14 @@ public class TestDocumentsAccess implements DocumentsAccess {
         lastUri.accept(docUri);
         return nextPath;
     }
+
+    public void assertCreatedDocument(DocumentInfo parent, String mimeType, String displayName) {
+        assertEquals(parent, mLastCreatedDoc.first);
+        assertEquals(mimeType, mLastCreatedDoc.second.mimeType);
+        assertEquals(displayName, mLastCreatedDoc.second.displayName);
+    }
+
+    public @Nullable Uri getLastCreatedDocumentUri() {
+        return mLastCreatedDoc.second.derivedUri;
+    }
 }
diff --git a/tests/common/com/android/documentsui/testing/TestEnv.java b/tests/common/com/android/documentsui/testing/TestEnv.java
index 9738d3cf9..54e91ca0b 100644
--- a/tests/common/com/android/documentsui/testing/TestEnv.java
+++ b/tests/common/com/android/documentsui/testing/TestEnv.java
@@ -61,6 +61,7 @@ public class TestEnv {
     public final TestRootsAccess roots = new TestRootsAccess();
     public final TestDocumentsAccess docs = new TestDocumentsAccess();
     public final TestFocusHandler focusHandler = new TestFocusHandler();
+    public final TestDialogController dialogs = new TestDialogController();
     public final TestModel model;
     public final TestModel archiveModel;
     public final SelectionManager selectionMgr;
@@ -84,7 +85,7 @@ public class TestEnv {
                 new TestActivityConfig(),
                 null,       //ScopedPreferences are not required for tests
                 null,   //a MessageBuilder is not required for tests
-                new TestDialogController(),
+                dialogs,
                 model);
         injector.selectionMgr = selectionMgr;
         injector.focusManager = new FocusManager(features, selectionMgr, null, null, 0);
@@ -158,17 +159,7 @@ public class TestEnv {
     }
 
     public void beforeAsserts() throws Exception {
-        // We need to wait on all AsyncTasks to finish AND to post results back.
-        // *** Results are posted on main thread ***, but tests run in their own
-        // thread. So even with our test executor we still have races.
-        //
-        // To work around this issue post our own runnable to the main thread
-        // which we presume will be the *last* runnable (after any from AsyncTasks)
-        // and then wait for our runnable to be called.
-        CountDownLatch latch = new CountDownLatch(1);
-        mExecutor.runAll();
-        new Handler(Looper.getMainLooper()).post(latch::countDown);
-        latch.await();
+        mExecutor.waitForTasks(30000); // 30 secs
     }
 
     public Executor lookupExecutor(String authority) {
diff --git a/tests/common/com/android/documentsui/testing/TestScheduledExecutorService.java b/tests/common/com/android/documentsui/testing/TestScheduledExecutorService.java
index bc05d4c85..bc008b005 100644
--- a/tests/common/com/android/documentsui/testing/TestScheduledExecutorService.java
+++ b/tests/common/com/android/documentsui/testing/TestScheduledExecutorService.java
@@ -17,13 +17,18 @@
 package com.android.documentsui.testing;
 
 import static junit.framework.Assert.assertFalse;
-import static junit.framework.Assert.fail;
+import static junit.framework.Assert.assertTrue;
+
+import android.os.Handler;
+import android.os.Looper;
+import android.os.SystemClock;
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Iterator;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Delayed;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -45,13 +50,11 @@ public class TestScheduledExecutorService implements ScheduledExecutorService {
     @Override
     public List<Runnable> shutdownNow() {
         this.shutdown = true;
-        return new ArrayList<>();
+        return Collections.emptyList();
     }
 
     public void assertShutdown() {
-        if (!shutdown) {
-            fail("Executor wasn't shut down.");
-        }
+        assertTrue("Executor wasn't shut down.", shutdown);
     }
 
     @Override
@@ -138,13 +141,9 @@ public class TestScheduledExecutorService implements ScheduledExecutorService {
     }
 
     public void runAll() {
-        final Iterator<TestFuture> iter = scheduled.iterator();
-        while (iter.hasNext()) {
-            TestFuture future = iter.next();
+        while (!scheduled.isEmpty()) {
+            TestFuture future = scheduled.remove(scheduled.size() - 1);
             future.runnable.run();
-
-            // Remove the job from scheduled after it finishes.
-            iter.remove();
         }
     }
 
@@ -158,6 +157,28 @@ public class TestScheduledExecutorService implements ScheduledExecutorService {
         assertFalse(isShutdown());
     }
 
+    public void waitForTasks(long millisTimeout) throws Exception {
+        millisTimeout = (millisTimeout > 0) ? millisTimeout : Long.MAX_VALUE;
+
+        final long startTime = SystemClock.uptimeMillis();
+
+        // We need to wait on all AsyncTasks to finish AND to post results back.
+        // *** Results are posted on main thread ***, but tests run in their own
+        // thread. So even with our test executor we still have races.
+        //
+        // To work around this issue post our own runnable to the main thread
+        // which we presume will be the *last* runnable (after any from AsyncTasks)
+        // and then wait for our runnable to be called.
+        while (!scheduled.isEmpty() && millisTimeout > 0) {
+            CountDownLatch latch = new CountDownLatch(1);
+            runAll();
+            new Handler(Looper.getMainLooper()).post(latch::countDown);
+            latch.await(millisTimeout, TimeUnit.MILLISECONDS);
+
+            millisTimeout -= (SystemClock.uptimeMillis() - startTime);
+        }
+    }
+
     static class TestFuture implements ScheduledFuture<Void> {
 
         final Runnable runnable;
diff --git a/tests/unit/com/android/documentsui/picker/ActionHandlerTest.java b/tests/unit/com/android/documentsui/picker/ActionHandlerTest.java
index bc6cd3c69..c65976cd9 100644
--- a/tests/unit/com/android/documentsui/picker/ActionHandlerTest.java
+++ b/tests/unit/com/android/documentsui/picker/ActionHandlerTest.java
@@ -16,17 +16,23 @@
 
 package com.android.documentsui.picker;
 
+import static junit.framework.Assert.assertTrue;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 
+import android.app.Activity;
+import android.content.ClipData;
 import android.content.Intent;
 import android.net.Uri;
+import android.os.AsyncTask;
 import android.provider.DocumentsContract;
 import android.provider.DocumentsContract.Path;
 import android.support.test.filters.MediumTest;
 import android.support.test.runner.AndroidJUnit4;
 
 import com.android.documentsui.R;
+import com.android.documentsui.base.DocumentStack;
 import com.android.documentsui.base.RootInfo;
 import com.android.documentsui.base.Shared;
 import com.android.documentsui.base.State;
@@ -35,8 +41,8 @@ import com.android.documentsui.testing.DocumentStackAsserts;
 import com.android.documentsui.testing.TestEnv;
 import com.android.documentsui.testing.TestRootsAccess;
 import com.android.documentsui.testing.TestLastAccessedStorage;
-import com.android.documentsui.ui.TestDialogController;
 
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -49,7 +55,6 @@ public class ActionHandlerTest {
 
     private TestEnv mEnv;
     private TestActivity mActivity;
-    private TestDialogController mDialogs;
     private ActionHandler<TestActivity> mHandler;
     private TestLastAccessedStorage mLastAccessed;
 
@@ -57,7 +62,6 @@ public class ActionHandlerTest {
     public void setUp() {
         mEnv = TestEnv.create();
         mActivity = TestActivity.create(mEnv);
-        mDialogs = new TestDialogController();
         mEnv.roots.configurePm(mActivity.packageMgr);
         mLastAccessed = new TestLastAccessedStorage();
 
@@ -72,9 +76,16 @@ public class ActionHandlerTest {
                 mLastAccessed
         );
 
-        mDialogs.confirmNext();
+        mEnv.dialogs.confirmNext();
 
         mEnv.selectionMgr.toggleSelection("1");
+
+        AsyncTask.setDefaultExecutor(mEnv.mExecutor);
+    }
+
+    @AfterClass
+    public static void tearDownOnce() {
+        AsyncTask.setDefaultExecutor(AsyncTask.THREAD_POOL_EXECUTOR);
     }
 
     @Test
@@ -123,30 +134,43 @@ public class ActionHandlerTest {
     }
 
     @Test
-    public void testOnLastAccessedStackLoaded_defaultToRecents_getContent() throws Exception {
-        testOnLastAccessedStackLoaded_defaultToRecentsOnAction(State.ACTION_GET_CONTENT);
+    public void testInitLocation_RestoresLastAccessedStack() throws Exception {
+        final DocumentStack stack =
+                new DocumentStack(TestRootsAccess.HAMMY, TestEnv.FOLDER_0, TestEnv.FOLDER_1);
+        mLastAccessed.setLastAccessed(mActivity, stack);
+
+        mHandler.initLocation(mActivity.getIntent());
+
+        mEnv.beforeAsserts();
+        assertEquals(stack, mEnv.state.stack);
+        mActivity.refreshCurrentRootAndDirectory.assertCalled();
     }
 
     @Test
-    public void testOnLastAccessedStackLoaded_defaultToRecents_open() throws Exception {
-        testOnLastAccessedStackLoaded_defaultToRecentsOnAction(State.ACTION_OPEN);
+    public void testInitLocation_DefaultToRecents_ActionGetContent() throws Exception {
+        testInitLocationDefaultToRecentsOnAction(State.ACTION_GET_CONTENT);
     }
 
     @Test
-    public void testOnLastAccessedStackLoaded_defaultToRecents_openTree() throws Exception {
-        testOnLastAccessedStackLoaded_defaultToRecentsOnAction(State.ACTION_OPEN_TREE);
+    public void testInitLocation_DefaultToRecents_ActionOpen() throws Exception {
+        testInitLocationDefaultToRecentsOnAction(State.ACTION_OPEN);
     }
 
     @Test
-    public void testOnLastAccessedStackLoaded_DefaultsToDownloads_create() throws Exception {
-        testOnLastAccessedStackLoaded_defaultToDownloadsOnAction(State.ACTION_CREATE);
+    public void testInitLocation_DefaultToRecents_ActionOpenTree() throws Exception {
+        testInitLocationDefaultToRecentsOnAction(State.ACTION_OPEN_TREE);
     }
 
     @Test
-    public void testOnLastAccessedStackLoaded_DefaultsToDownloads_pickCopyDestination()
-            throws Exception {
-        testOnLastAccessedStackLoaded_defaultToDownloadsOnAction(
-                State.ACTION_PICK_COPY_DESTINATION);
+    public void testInitLocation_DefaultsToDownloads_ActionCreate() throws Exception {
+        mEnv.state.action = State.ACTION_CREATE;
+        mActivity.resources.bools.put(R.bool.show_documents_root, false);
+
+        mActivity.refreshCurrentRootAndDirectory.assertNotCalled();
+
+        mHandler.initLocation(mActivity.getIntent());
+
+        assertRootPicked(TestRootsAccess.DOWNLOADS.getUri());
     }
 
     @Test
@@ -158,24 +182,241 @@ public class ActionHandlerTest {
         mActivity.refreshCurrentRootAndDirectory.assertCalled();
     }
 
-    private void testOnLastAccessedStackLoaded_defaultToRecentsOnAction(@ActionType int action) {
-        mEnv.state.action = action;
-        mActivity.refreshCurrentRootAndDirectory.assertNotCalled();
+    @Test
+    public void testPickDocument_SetsCorrectResultAndFinishes_ActionPickCopyDestination()
+            throws Exception {
 
-        mHandler.onLastAccessedStackLoaded(null);
+        mEnv.state.action = State.ACTION_PICK_COPY_DESTINATION;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+        mEnv.state.stack.push(TestEnv.FOLDER_2);
 
-        assertEquals(TestRootsAccess.RECENTS, mEnv.state.stack.getRoot());
-        mActivity.refreshCurrentRootAndDirectory.assertCalled();
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.pickDocument(TestEnv.FOLDER_2);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FOLDER_2.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testPickDocument_SetsCorrectResultAndFinishes_ActionOpenTree() throws Exception {
+        mEnv.state.action = State.ACTION_OPEN_TREE;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+        mEnv.state.stack.push(TestEnv.FOLDER_2);
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.pickDocument(TestEnv.FOLDER_2);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, true);
+        assertContent(result, DocumentsContract.buildTreeDocumentUri(
+                TestRootsAccess.HOME.authority, TestEnv.FOLDER_2.documentId));
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testSaveDocument_SetsCorrectResultAndFinishes() throws Exception {
+        mEnv.state.action = State.ACTION_CREATE;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+
+        final String mimeType = "audio/aac";
+        final String displayName = "foobar.m4a";
+
+        mHandler.saveDocument(mimeType, displayName, (boolean inProgress) -> {});
+
+        mEnv.beforeAsserts();
+
+        mEnv.docs.assertCreatedDocument(TestEnv.FOLDER_1, mimeType, displayName);
+        final Uri docUri = mEnv.docs.getLastCreatedDocumentUri();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, docUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testSaveDocument_ConfirmsOverwrite() {
+        mEnv.state.action = State.ACTION_CREATE;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+
+        mHandler.saveDocument(null, TestEnv.FILE_JPG);
+
+        mEnv.dialogs.assertOverwriteConfirmed(TestEnv.FILE_JPG);
+    }
+
+    @Test
+    public void testFinishPicking_SetsCorrectResultAndFinishes_ActionGetContent() throws Exception {
+        mEnv.state.action = State.ACTION_GET_CONTENT;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.finishPicking(TestEnv.FILE_JPG.derivedUri);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FILE_JPG.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testFinishPicking_SetsCorrectResultAndFinishes_ActionGetContent_MultipleSelection()
+            throws Exception {
+        mEnv.state.action = State.ACTION_GET_CONTENT;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+        mEnv.state.acceptMimes = new String[] { "image/*" };
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.finishPicking(TestEnv.FILE_JPG.derivedUri, TestEnv.FILE_GIF.derivedUri);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, false);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FILE_JPG.derivedUri, TestEnv.FILE_GIF.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testFinishPicking_SetsCorrectResultAndFinishes_ActionOpen() throws Exception {
+        mEnv.state.action = State.ACTION_OPEN;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.finishPicking(TestEnv.FILE_JPG.derivedUri);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FILE_JPG.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testFinishPicking_SetsCorrectResultAndFinishes_ActionOpen_MultipleSelection()
+            throws Exception {
+        mEnv.state.action = State.ACTION_OPEN;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+        mEnv.state.acceptMimes = new String[] { "image/*" };
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.finishPicking(TestEnv.FILE_JPG.derivedUri, TestEnv.FILE_GIF.derivedUri);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FILE_JPG.derivedUri, TestEnv.FILE_GIF.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
+    }
+
+    @Test
+    public void testFinishPicking_SetsCorrectResultAndFinishes_ActionCreate() throws Exception {
+        mEnv.state.action = State.ACTION_CREATE;
+        mEnv.state.stack.changeRoot(TestRootsAccess.HOME);
+        mEnv.state.stack.push(TestEnv.FOLDER_1);
+
+        mActivity.finishedHandler.assertNotCalled();
+
+        mHandler.finishPicking(TestEnv.FILE_JPG.derivedUri);
+
+        mEnv.beforeAsserts();
+
+        assertLastAccessedStackUpdated();
+
+        assertEquals(Activity.RESULT_OK, (long) mActivity.setResult.getLastValue().first);
+        final Intent result = mActivity.setResult.getLastValue().second;
+        assertPermission(result, Intent.FLAG_GRANT_READ_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_WRITE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION, true);
+        assertPermission(result, Intent.FLAG_GRANT_PREFIX_URI_PERMISSION, false);
+        assertContent(result, TestEnv.FILE_JPG.derivedUri);
+
+        mActivity.finishedHandler.assertCalled();
     }
 
-    private void testOnLastAccessedStackLoaded_defaultToDownloadsOnAction(@ActionType int action)
+    private void testInitLocationDefaultToRecentsOnAction(@ActionType int action)
             throws Exception {
         mEnv.state.action = action;
+
         mActivity.refreshCurrentRootAndDirectory.assertNotCalled();
 
-        mHandler.onLastAccessedStackLoaded(null);
+        mHandler.initLocation(mActivity.getIntent());
 
-        assertRootPicked(TestRootsAccess.DOWNLOADS.getUri());
+        mEnv.beforeAsserts();
+        assertEquals(TestRootsAccess.RECENTS, mEnv.state.stack.getRoot());
+        mActivity.refreshCurrentRootAndDirectory.assertCalled();
     }
 
     private void assertRootPicked(Uri expectedUri) throws Exception {
@@ -186,4 +427,35 @@ public class ActionHandlerTest {
         assertNotNull(root);
         assertEquals(expectedUri, root.getUri());
     }
+
+    private void assertLastAccessedStackUpdated() {
+        assertEquals(
+                mEnv.state.stack, mLastAccessed.getLastAccessed(mActivity, mEnv.roots, mEnv.state));
+    }
+
+    private void assertPermission(Intent intent, int permission, boolean granted) {
+        int flags = intent.getFlags();
+
+        if (granted) {
+            assertEquals(permission, flags & permission);
+        } else {
+            assertEquals(0, flags & permission);
+        }
+    }
+
+    private void assertContent(Intent intent, Uri... contents) {
+        if (contents.length == 1) {
+            assertEquals(contents[0], intent.getData());
+        } else {
+            ClipData clipData = intent.getClipData();
+
+            assertNotNull(clipData);
+            for (int i = 0; i < mEnv.state.acceptMimes.length; ++i) {
+                assertEquals(mEnv.state.acceptMimes[i], clipData.getDescription().getMimeType(i));
+            }
+            for (int i = 0; i < contents.length; ++i) {
+                assertEquals(contents[i], clipData.getItemAt(i).getUri());
+            }
+        }
+    }
 }
diff --git a/tests/unit/com/android/documentsui/picker/TestActivity.java b/tests/unit/com/android/documentsui/picker/TestActivity.java
index b1622ae16..8ce91b80e 100644
--- a/tests/unit/com/android/documentsui/picker/TestActivity.java
+++ b/tests/unit/com/android/documentsui/picker/TestActivity.java
@@ -16,17 +16,35 @@
 
 package com.android.documentsui.picker;
 
+import android.content.Intent;
+import android.util.Pair;
+
 import com.android.documentsui.testing.TestEnv;
+import com.android.documentsui.testing.TestEventListener;
 
 import org.mockito.Mockito;
 
 public abstract class TestActivity extends AbstractBase {
 
+    public TestEventListener<Pair<Integer, Intent>> setResult;
+
     public static TestActivity create(TestEnv env) {
         TestActivity activity = Mockito.mock(TestActivity.class, Mockito.CALLS_REAL_METHODS);
         activity.init(env);
         return activity;
     }
+
+    @Override
+    public void init(TestEnv env) {
+        super.init(env);
+
+        setResult = new TestEventListener<>();
+    }
+
+    @Override
+    public void setResult(int resultCode, Intent intent, int notUsed) {
+        setResult.accept(Pair.create(resultCode, intent));
+    }
 }
 
 // Trick Mockito into finding our Addons methods correctly. W/o this