diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
index 322246a..cb49535e 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
@@ -56,22 +56,26 @@
         mDatabase = database;
     }
 
-    synchronized String putDeviceDocument(int deviceId, String name, MtpRoot[] roots) {
+    /**
+     * Puts device information to database.
+     * @return If device is added to the database.
+     */
+    synchronized boolean putDeviceDocument(MtpDeviceRecord device) {
         final SQLiteDatabase database = mDatabase.getSQLiteDatabase();
         Preconditions.checkState(mMappingMode.containsKey(/* no parent for root */ null));
         database.beginTransaction();
         try {
             final ContentValues[] valuesList = new ContentValues[1];
             valuesList[0] = new ContentValues();
-            MtpDatabase.getDeviceDocumentValues(valuesList[0], deviceId, name, roots);
-            putDocuments(
+            MtpDatabase.getDeviceDocumentValues(valuesList[0], device);
+            final boolean changed = putDocuments(
                     valuesList,
                     COLUMN_PARENT_DOCUMENT_ID + " IS NULL",
                     EMPTY_ARGS,
                     /* heuristic */ false,
                     COLUMN_DEVICE_ID);
             database.setTransactionSuccessful();
-            return valuesList[0].getAsString(Document.COLUMN_DOCUMENT_ID);
+            return changed;
         } finally {
             database.endTransaction();
         }
@@ -249,7 +253,7 @@
      * If the mapping mode is not heuristic, it just adds the rows to the database or updates the
      * existing rows with the new values. If the mapping mode is heuristic, it adds some new rows as
      * 'pending' state when that rows may be corresponding to existing 'invalidated' rows. Then
-     * {@link #stopAddingDocuments(String, String[], String)} turns the pending rows into 'valid'
+     * {@link #stopAddingDocuments(String)} turns the pending rows into 'valid'
      * rows. If the methods adds rows to database, it updates valueList with correct document ID.
      *
      * @param valuesList Values for documents to be stored in the database.
@@ -452,12 +456,15 @@
         final SQLiteDatabase database = mDatabase.getSQLiteDatabase();
         values.clear();
         final Cursor cursor = database.query(table, null, selection, args, null, null, null, "1");
-        if (cursor.getCount() == 0) {
-            return;
+        try {
+            if (cursor.getCount() == 0) {
+                return;
+            }
+            cursor.moveToNext();
+            DatabaseUtils.cursorRowToContentValues(cursor, values);
+        } finally {
+            cursor.close();
         }
-        cursor.moveToNext();
-        DatabaseUtils.cursorRowToContentValues(cursor, values);
-        cursor.close();
     }
 
     /**
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
index 1ba3e31..10941eb 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
@@ -18,11 +18,11 @@
 
 import static com.android.mtp.MtpDatabaseConstants.*;
 
+import android.annotation.Nullable;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.res.Resources;
 import android.database.Cursor;
-import android.database.MatrixCursor;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteOpenHelper;
 import android.database.sqlite.SQLiteQueryBuilder;
@@ -36,8 +36,6 @@
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.FileNotFoundException;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -183,6 +181,27 @@
         deleteDocumentsAndRoots(COLUMN_DEVICE_ID + "=?", strings(deviceId));
     }
 
+    @Nullable String getDocumentIdForDevice(int deviceId) {
+        final Cursor cursor = mDatabase.query(
+                TABLE_DOCUMENTS,
+                strings(Document.COLUMN_DOCUMENT_ID),
+                COLUMN_DOCUMENT_TYPE + " = ? AND " + COLUMN_DEVICE_ID + " = ?",
+                strings(DOCUMENT_TYPE_DEVICE, deviceId),
+                null,
+                null,
+                null,
+                "1");
+        try {
+            if (cursor.moveToNext()) {
+                return cursor.getString(0);
+            } else {
+                return null;
+            }
+        } finally {
+            cursor.close();
+        }
+    }
+
     /**
      * Obtains parent document ID.
      * @param documentId
@@ -361,23 +380,22 @@
         context.deleteDatabase(DATABASE_NAME);
     }
 
-    static void getDeviceDocumentValues(
-            ContentValues values, int deviceId, String name, MtpRoot[] roots) {
+    static void getDeviceDocumentValues(ContentValues values, MtpDeviceRecord device) {
         values.clear();
-        values.put(COLUMN_DEVICE_ID, deviceId);
+        values.put(COLUMN_DEVICE_ID, device.deviceId);
         values.putNull(COLUMN_STORAGE_ID);
         values.putNull(COLUMN_OBJECT_HANDLE);
         values.putNull(COLUMN_PARENT_DOCUMENT_ID);
         values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
         values.put(COLUMN_DOCUMENT_TYPE, DOCUMENT_TYPE_DEVICE);
         values.put(Document.COLUMN_MIME_TYPE, Document.MIME_TYPE_DIR);
-        values.put(Document.COLUMN_DISPLAY_NAME, name);
+        values.put(Document.COLUMN_DISPLAY_NAME, device.name);
         values.putNull(Document.COLUMN_SUMMARY);
         values.putNull(Document.COLUMN_LAST_MODIFIED);
         values.put(Document.COLUMN_ICON, R.drawable.ic_root_mtp);
         values.put(Document.COLUMN_FLAGS, 0);
         long size = 0;
-        for (final MtpRoot root : roots) {
+        for (final MtpRoot root : device.roots) {
             size += root.mMaxCapacity - root.mFreeSpace;
         }
         values.put(Document.COLUMN_SIZE, size);
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDeviceRecord.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDeviceRecord.java
new file mode 100644
index 0000000..71df5c1
--- /dev/null
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDeviceRecord.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2015 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.mtp;
+
+class MtpDeviceRecord {
+    public final int deviceId;
+    public final String name;
+    public final boolean opened;
+    public final MtpRoot[] roots;
+
+    MtpDeviceRecord(int deviceId, String name, boolean opened, MtpRoot[] roots) {
+        this.deviceId = deviceId;
+        this.name = name;
+        this.opened = opened;
+        this.roots = roots;
+    }
+}
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
index caa2024..57a68ba 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
@@ -247,7 +247,7 @@
             closeDeviceInternal(deviceId);
             mDatabase.removeDeviceRows(deviceId);
         }
-        mRootScanner.notifyChange();
+        mRootScanner.resume();
     }
 
     int[] getOpenedDeviceIds() {
@@ -258,7 +258,12 @@
 
     String getDeviceName(int deviceId) throws IOException {
         synchronized (mDeviceListLock) {
-            return mMtpManager.getDeviceName(deviceId);
+            for (final MtpDeviceRecord device : mMtpManager.getDevices()) {
+                if (device.deviceId == deviceId) {
+                    return device.name;
+                }
+            }
+            throw new IOException("Not found the device: " + Integer.toString(deviceId));
         }
     }
 
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpManager.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpManager.java
index e7f94b7..88cab8b 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpManager.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpManager.java
@@ -17,8 +17,10 @@
 package com.android.mtp;
 
 import android.content.Context;
+import android.hardware.usb.UsbConstants;
 import android.hardware.usb.UsbDevice;
 import android.hardware.usb.UsbDeviceConnection;
+import android.hardware.usb.UsbInterface;
 import android.hardware.usb.UsbManager;
 import android.mtp.MtpConstants;
 import android.mtp.MtpDevice;
@@ -26,12 +28,14 @@
 import android.mtp.MtpObjectInfo;
 import android.os.CancellationSignal;
 import android.os.ParcelFileDescriptor;
+import android.util.Log;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.util.ArrayList;
 
 /**
  * The model wrapping android.mtp API.
@@ -39,6 +43,27 @@
 class MtpManager {
     final static int OBJECT_HANDLE_ROOT_CHILDREN = -1;
 
+    /**
+     * Subclass for PTP.
+     */
+    private static final int SUBCLASS_STILL_IMAGE_CAPTURE = 1;
+
+    /**
+     * Subclass for Android style MTP.
+     */
+    private static final int SUBCLASS_MTP = 0xff;
+
+    /**
+     * Protocol for Picture Transfer Protocol (PIMA 15470).
+     */
+    private static final int PROTOCOL_PICTURE_TRANSFER = 1;
+
+    /**
+     * Protocol for Android style MTP.
+     */
+    private static final int PROTOCOL_MTP = 0;
+
+
     private final UsbManager mManager;
     // TODO: Save and restore the set of opened device.
     private final SparseArray<MtpDevice> mDevices = new SparseArray<>();
@@ -92,6 +117,33 @@
         mDevices.remove(deviceId);
     }
 
+    synchronized MtpDeviceRecord[] getDevices() {
+        final ArrayList<MtpDeviceRecord> devices = new ArrayList<>();
+        for (UsbDevice device : mManager.getDeviceList().values()) {
+            if (!isMtpDevice(device)) {
+                continue;
+            }
+            final boolean opened = mDevices.get(device.getDeviceId()) != null;
+            final String name = device.getProductName();
+            MtpRoot[] roots;
+            if (opened) {
+                try {
+                    roots = getRoots(device.getDeviceId());
+                } catch (IOException exp) {
+                    Log.e(MtpDocumentsProvider.TAG, exp.getMessage());
+                    // If we failed to fetch roots for the device, we still returns device model
+                    // with an empty set of roots so that the device is shown DocumentsUI as long as
+                    // the device is physically connected.
+                    roots = new MtpRoot[0];
+                }
+            } else {
+                roots = new MtpRoot[0];
+            }
+            devices.add(new MtpDeviceRecord(device.getDeviceId(), name, opened, roots));
+        }
+        return devices.toArray(new MtpDeviceRecord[devices.size()]);
+    }
+
     synchronized int[] getOpenedDeviceIds() {
         final int[] result = new int[mDevices.size()];
         for (int i = 0; i < result.length; i++) {
@@ -100,28 +152,6 @@
         return result;
     }
 
-    String getDeviceName(int deviceId) throws IOException {
-        return getDevice(deviceId).getDeviceInfo().getModel();
-    }
-
-    MtpRoot[] getRoots(int deviceId) throws IOException {
-        final MtpDevice device = getDevice(deviceId);
-        synchronized (device) {
-            final int[] storageIds = device.getStorageIds();
-            if (storageIds == null) {
-                throw new IOException("Failed to obtain storage IDs.");
-            }
-            final MtpRoot[] results = new MtpRoot[storageIds.length];
-            for (int i = 0; i < storageIds.length; i++) {
-                results[i] = new MtpRoot(
-                        device.getDeviceId(),
-                        device.getDeviceInfo().getModel(),
-                        device.getStorageInfo(storageIds[i]));
-            }
-            return results;
-        }
-    }
-
     MtpObjectInfo getObjectInfo(int deviceId, int objectHandle)
             throws IOException {
         final MtpDevice device = getDevice(deviceId);
@@ -212,4 +242,40 @@
         }
         return device;
     }
+
+    private MtpRoot[] getRoots(int deviceId) throws IOException {
+        final MtpDevice device = getDevice(deviceId);
+        synchronized (device) {
+            final int[] storageIds = device.getStorageIds();
+            if (storageIds == null) {
+                throw new IOException("Failed to obtain storage IDs.");
+            }
+            final MtpRoot[] results = new MtpRoot[storageIds.length];
+            for (int i = 0; i < storageIds.length; i++) {
+                results[i] = new MtpRoot(
+                        device.getDeviceId(),
+                        device.getDeviceInfo().getModel(),
+                        device.getStorageInfo(storageIds[i]));
+            }
+            return results;
+        }
+    }
+
+    static boolean isMtpDevice(UsbDevice device) {
+        for (int i = 0; i < device.getInterfaceCount(); i++) {
+            final UsbInterface usbInterface = device.getInterface(i);
+            if ((usbInterface.getInterfaceClass() == UsbConstants.USB_CLASS_STILL_IMAGE &&
+                    usbInterface.getInterfaceSubclass() == SUBCLASS_STILL_IMAGE_CAPTURE &&
+                    usbInterface.getInterfaceProtocol() == PROTOCOL_PICTURE_TRANSFER)) {
+                return true;
+            }
+            if (usbInterface.getInterfaceClass() == UsbConstants.USB_SUBCLASS_VENDOR_SPEC &&
+                    usbInterface.getInterfaceSubclass() == SUBCLASS_MTP &&
+                    usbInterface.getInterfaceProtocol() == PROTOCOL_MTP &&
+                    usbInterface.getName().equals("MTP")) {
+                return true;
+            }
+        }
+        return false;
+    }
 }
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
index 197dec8..e6c2726 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
@@ -2,12 +2,10 @@
 
 import android.content.ContentResolver;
 import android.content.res.Resources;
-import android.database.sqlite.SQLiteException;
 import android.net.Uri;
 import android.os.Process;
 import android.provider.DocumentsContract;
 import android.util.Log;
-import android.util.SparseArray;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -108,36 +106,29 @@
             Process.setThreadPriority(Process.THREAD_PRIORITY_BACKGROUND);
             int pollingCount = 0;
             while (!Thread.interrupted()) {
-                final int[] deviceIds = mManager.getOpenedDeviceIds();
-                final Map<String, MtpRoot[]> rootsMap = new HashMap<>();
                 boolean changed = false;
 
                 // Update devices.
+                final MtpDeviceRecord[] devices = mManager.getDevices();
                 mDatabase.getMapper().startAddingDocuments(null /* parentDocumentId */);
-                for (final int deviceId : deviceIds) {
-                    try {
-                        final MtpRoot[] roots = mManager.getRoots(deviceId);
-                        final String id = mDatabase.getMapper().putDeviceDocument(
-                                deviceId,
-                                mManager.getDeviceName(deviceId),
-                                roots);
-                        if (id != null) {
-                            changed = true;
-                            rootsMap.put(id, roots);
-                        }
-                    } catch (IOException exception) {
-                        // The error may happen on the device. We would like to continue getting
-                        // roots for other devices.
-                        Log.e(MtpDocumentsProvider.TAG, exception.getMessage());
+                for (final MtpDeviceRecord device : devices) {
+                    if (mDatabase.getMapper().putDeviceDocument(device)) {
+                        changed = true;
                     }
                 }
-                mDatabase.getMapper().stopAddingDocuments(null /* parentDocumentId */);
+                if (mDatabase.getMapper().stopAddingDocuments(null /* parentDocumentId */)) {
+                    changed = true;
+                }
 
                 // Update roots.
-                for (final String documentId : rootsMap.keySet()) {
+                for (final MtpDeviceRecord device : devices) {
+                    final String documentId = mDatabase.getDocumentIdForDevice(device.deviceId);
+                    if (documentId == null) {
+                        continue;
+                    }
                     mDatabase.getMapper().startAddingDocuments(documentId);
                     if (mDatabase.getMapper().putRootDocuments(
-                            documentId, mResources, rootsMap.get(documentId))) {
+                            documentId, mResources, device.roots)) {
                         changed = true;
                     }
                     if (mDatabase.getMapper().stopAddingDocuments(documentId)) {
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
index 6d57c5b..b745175 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
@@ -871,4 +871,12 @@
             cursor.close();
         }
     }
+
+    public void testGetDocumentIdForDevice() {
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(100, "Device", true, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+        assertEquals("1", mDatabase.getDocumentIdForDevice(100));
+    }
 }
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
index b0e9722..884b1e2 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
@@ -23,7 +23,6 @@
 import android.provider.DocumentsContract.Root;
 import android.provider.DocumentsContract;
 import android.test.AndroidTestCase;
-import android.test.suitebuilder.annotation.MediumTest;
 import android.test.suitebuilder.annotation.SmallTest;
 
 import java.io.FileNotFoundException;
@@ -56,17 +55,20 @@
 
     public void testOpenAndCloseDevice() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mMtpManager.setRoots(0, new MtpRoot[] {
-                new MtpRoot(
-                        0 /* deviceId */,
-                        1 /* storageId */,
-                        "Device A" /* device model name */,
-                        "Storage A" /* volume description */,
-                        1024 /* free space */,
-                        2048 /* total space */,
-                        "" /* no volume identifier */)
-        });
+        mMtpManager.addValidDevice(new MtpDeviceRecord(
+                0,
+                "Device",
+                false /* unopened */,
+                new MtpRoot[] {
+                    new MtpRoot(
+                            0 /* deviceId */,
+                            1 /* storageId */,
+                            "Device A" /* device model name */,
+                            "Storage A" /* volume description */,
+                            1024 /* free space */,
+                            2048 /* total space */,
+                            "" /* no volume identifier */)
+                }));
 
         mProvider.openDevice(0);
         mResolver.waitForNotification(ROOTS_URI, 1);
@@ -92,45 +94,54 @@
         }
 
         // Check if the following notification is the first one or not.
-        mMtpManager.addValidDevice(0);
-        mMtpManager.setRoots(0, new MtpRoot[] {
-                new MtpRoot(
-                        0 /* deviceId */,
-                        1 /* storageId */,
-                        "Device A" /* device model name */,
-                        "Storage A" /* volume description */,
-                        1024 /* free space */,
-                        2048 /* total space */,
-                        "" /* no volume identifier */)
-        });
+        mMtpManager.addValidDevice(new MtpDeviceRecord(
+                0,
+                "Device",
+                false /* unopened */,
+                new MtpRoot[] {
+                    new MtpRoot(
+                            0 /* deviceId */,
+                            1 /* storageId */,
+                            "Device A" /* device model name */,
+                            "Storage A" /* volume description */,
+                            1024 /* free space */,
+                            2048 /* total space */,
+                            "" /* no volume identifier */)
+                }));
         mProvider.openDevice(0);
         mResolver.waitForNotification(ROOTS_URI, 1);
     }
 
     public void testQueryRoots() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mMtpManager.addValidDevice(1);
-        mMtpManager.setRoots(0, new MtpRoot[] {
-                new MtpRoot(
-                        0 /* deviceId */,
-                        1 /* storageId */,
-                        "Device A" /* device model name */,
-                        "Storage A" /* volume description */,
-                        1024 /* free space */,
-                        2048 /* total space */,
-                        "" /* no volume identifier */)
-        });
-        mMtpManager.setRoots(1, new MtpRoot[] {
-                new MtpRoot(
-                        1 /* deviceId */,
-                        1 /* storageId */,
-                        "Device B" /* device model name */,
-                        "Storage B" /* volume description */,
-                        2048 /* free space */,
-                        4096 /* total space */,
-                        "Identifier B" /* no volume identifier */)
-        });
+        mMtpManager.addValidDevice(new MtpDeviceRecord(
+                0,
+                "Device",
+                false /* unopened */,
+                new MtpRoot[] {
+                        new MtpRoot(
+                                0 /* deviceId */,
+                                1 /* storageId */,
+                                "Device A" /* device model name */,
+                                "Storage A" /* volume description */,
+                                1024 /* free space */,
+                                2048 /* total space */,
+                                "" /* no volume identifier */)
+                }));
+        mMtpManager.addValidDevice(new MtpDeviceRecord(
+                1,
+                "Device",
+                false /* unopened */,
+                new MtpRoot[] {
+                    new MtpRoot(
+                            1 /* deviceId */,
+                            1 /* storageId */,
+                            "Device B" /* device model name */,
+                            "Storage B" /* volume description */,
+                            2048 /* free space */,
+                            4096 /* total space */,
+                            "Identifier B" /* no volume identifier */)
+                }));
 
         {
             mProvider.openDevice(0);
@@ -138,11 +149,11 @@
             final Cursor cursor = mProvider.queryRoots(null);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals("2", cursor.getString(0));
+            assertEquals("3", cursor.getString(0));
             assertEquals(Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE, cursor.getInt(1));
             assertEquals(R.drawable.ic_root_mtp, cursor.getInt(2));
             assertEquals("Device A Storage A", cursor.getString(3));
-            assertEquals("2", cursor.getString(4));
+            assertEquals("3", cursor.getString(4));
             assertEquals(1024, cursor.getInt(5));
         }
 
@@ -164,19 +175,22 @@
 
     public void testQueryRoots_error() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mMtpManager.addValidDevice(1);
-        // Not set roots for device 0 so that MtpManagerMock#getRoots throws IOException.
-        mMtpManager.setRoots(1, new MtpRoot[] {
-                new MtpRoot(
-                        1 /* deviceId */,
-                        1 /* storageId */,
-                        "Device B" /* device model name */,
-                        "Storage B" /* volume description */,
-                        2048 /* free space */,
-                        4096 /* total space */,
-                        "Identifier B" /* no volume identifier */)
-        });
+        mMtpManager.addValidDevice(
+                new MtpDeviceRecord(0, "Device", false /* unopened */, new MtpRoot[0]));
+        mMtpManager.addValidDevice(new MtpDeviceRecord(
+                1,
+                "Device",
+                false /* unopened */,
+                new MtpRoot[] {
+                    new MtpRoot(
+                            1 /* deviceId */,
+                            1 /* storageId */,
+                            "Device B" /* device model name */,
+                            "Storage B" /* volume description */,
+                            2048 /* free space */,
+                            4096 /* total space */,
+                            "Identifier B" /* no volume identifier */)
+                }));
         {
             mProvider.openDevice(0);
             mProvider.openDevice(1);
@@ -185,20 +199,17 @@
             final Cursor cursor = mProvider.queryRoots(null);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals("2", cursor.getString(0));
+            assertEquals("3", cursor.getString(0));
             assertEquals(Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE, cursor.getInt(1));
             assertEquals(R.drawable.ic_root_mtp, cursor.getInt(2));
             assertEquals("Device B Storage B", cursor.getString(3));
-            assertEquals("2", cursor.getString(4));
+            assertEquals("3", cursor.getString(4));
             assertEquals(2048, cursor.getInt(5));
         }
     }
 
     public void testQueryDocument() throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
-
         setupRoots(0, new MtpRoot[] { new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "") });
         setupDocuments(
                 0,
@@ -236,9 +247,6 @@
     public void testQueryDocument_directory()
             throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
-
         setupRoots(0, new MtpRoot[] { new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "") });
         setupDocuments(
                 0,
@@ -274,9 +282,6 @@
     public void testQueryDocument_forRoot()
             throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
-
         setupRoots(0, new MtpRoot[] {
                 new MtpRoot(
                         0 /* deviceId */,
@@ -301,9 +306,6 @@
 
     public void testQueryChildDocuments() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
-
         setupRoots(0, new MtpRoot[] { new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "") });
         setupDocuments(
                 0,
@@ -337,8 +339,6 @@
 
     public void testQueryChildDocuments_cursorError() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
         try {
             mProvider.queryChildDocuments("1", null, null);
             fail();
@@ -349,8 +349,6 @@
 
     public void testQueryChildDocuments_documentError() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
         setupRoots(0, new MtpRoot[] { new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "") });
         mMtpManager.setObjectHandles(0, 0, -1, new int[] { 1 });
         try {
@@ -363,8 +361,6 @@
 
     public void testDeleteDocument() throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
         setupRoots(0, new MtpRoot[] {
                 new MtpRoot(0, 0, "Device", "Storage", 0, 0, "")
         });
@@ -385,8 +381,6 @@
     public void testDeleteDocument_error()
             throws IOException, InterruptedException, TimeoutException {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mMtpManager.addValidDevice(0);
-        mProvider.openDevice(0);
         setupRoots(0, new MtpRoot[] {
                 new MtpRoot(0, 0, "Device", "Storage", 0, 0, "")
         });
@@ -427,9 +421,11 @@
     }
 
     private String[] setupRoots(int deviceId, MtpRoot[] roots)
-            throws FileNotFoundException, InterruptedException, TimeoutException {
+            throws InterruptedException, TimeoutException, IOException {
         final int changeCount = mResolver.getChangeCount(ROOTS_URI);
-        mMtpManager.setRoots(deviceId, roots);
+        mMtpManager.addValidDevice(
+                new MtpDeviceRecord(deviceId, "Device", false /* unopened */, roots));
+        mProvider.openDevice(deviceId);
         mResolver.waitForNotification(ROOTS_URI, changeCount + 1);
         return getStrings(mProvider.queryRoots(strings(DocumentsContract.Root.COLUMN_ROOT_ID)));
     }
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
index ddc18a4..bbd0a30 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
@@ -19,14 +19,12 @@
 import android.content.Context;
 import android.mtp.MtpObjectInfo;
 import android.os.ParcelFileDescriptor;
+import android.util.SparseArray;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
-import java.util.TreeSet;
 
 public class TestMtpManager extends MtpManager {
     public static final int CREATED_DOCUMENT_HANDLE = 1000;
@@ -35,9 +33,7 @@
         return Arrays.toString(args);
     }
 
-    private final Set<Integer> mValidDevices = new HashSet<>();
-    private final Set<Integer> mOpenedDevices = new TreeSet<>();
-    private final Map<Integer, MtpRoot[]> mRoots = new HashMap<>();
+    private final SparseArray<MtpDeviceRecord> mDevices = new SparseArray<>();
     private final Map<String, MtpObjectInfo> mObjectInfos = new HashMap<>();
     private final Map<String, int[]> mObjectHandles = new HashMap<>();
     private final Map<String, byte[]> mThumbnailBytes = new HashMap<>();
@@ -47,18 +43,14 @@
         super(context);
     }
 
-    void addValidDevice(int deviceId) {
-        mValidDevices.add(deviceId);
+    void addValidDevice(MtpDeviceRecord device) {
+        mDevices.put(device.deviceId, device);
     }
 
     void setObjectHandles(int deviceId, int storageId, int parentHandle, int[] objectHandles) {
         mObjectHandles.put(pack(deviceId, storageId, parentHandle), objectHandles);
     }
 
-    void setRoots(int deviceId, MtpRoot[] roots) {
-        mRoots.put(deviceId, roots);
-    }
-
     void setObjectInfo(int deviceId, MtpObjectInfo objectInfo) {
         mObjectInfos.put(pack(deviceId, objectInfo.getObjectHandle()), objectInfo);
     }
@@ -76,28 +68,40 @@
     }
 
     @Override
+    MtpDeviceRecord[] getDevices() {
+        final MtpDeviceRecord[] result = new MtpDeviceRecord[mDevices.size()];
+        for (int i = 0; i < mDevices.size(); i++) {
+            final MtpDeviceRecord device = mDevices.valueAt(i);
+            if (device.opened) {
+                result[i] = device;
+            } else {
+                result[i] = new MtpDeviceRecord(
+                        device.deviceId, device.name, device.opened, new MtpRoot[0]);
+            }
+        }
+        return result;
+    }
+
+    @Override
     void openDevice(int deviceId) throws IOException {
-        if (!mValidDevices.contains(deviceId) || mOpenedDevices.contains(deviceId)) {
+        final MtpDeviceRecord device = mDevices.get(deviceId);
+        if (device == null || device.opened) {
             throw new IOException();
         }
-        mOpenedDevices.add(deviceId);
+        mDevices.put(
+                deviceId,
+                new MtpDeviceRecord(device.deviceId, device.name, true, device.roots));
     }
 
     @Override
     void closeDevice(int deviceId) throws IOException {
-        if (!mValidDevices.contains(deviceId) || !mOpenedDevices.contains(deviceId)) {
+        final MtpDeviceRecord device = mDevices.get(deviceId);
+        if (device == null || !device.opened) {
             throw new IOException();
         }
-        mOpenedDevices.remove(deviceId);
-    }
-
-    @Override
-    MtpRoot[] getRoots(int deviceId) throws IOException {
-        if (mRoots.containsKey(deviceId)) {
-            return mRoots.get(deviceId);
-        } else {
-            throw new IOException("getRoots error: " + Integer.toString(deviceId));
-        }
+        mDevices.put(
+                deviceId,
+                new MtpDeviceRecord(device.deviceId, device.name, false, device.roots));
     }
 
     @Override
@@ -189,16 +193,14 @@
 
     @Override
     int[] getOpenedDeviceIds() {
-        int i = 0;
-        final int[] result = new int[mOpenedDevices.size()];
-        for (int deviceId : mOpenedDevices) {
-            result[i++] = deviceId;
+        final int[] result = new int[mDevices.size()];
+        int count = 0;
+        for (int i = 0; i < mDevices.size(); i++) {
+            final MtpDeviceRecord device = mDevices.valueAt(i);
+            if (device.opened) {
+                result[count++] = device.deviceId;
+            }
         }
-        return result;
-    }
-
-    @Override
-    String getDeviceName(int deviceId) throws IOException {
-        return "Device";
+        return Arrays.copyOf(result, count);
     }
 }
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestUtil.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestUtil.java
index f910321..611e831 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestUtil.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestUtil.java
@@ -120,9 +120,19 @@
     private static void waitForStorages(
             TestResultInstrumentation instrumentation,
             MtpManager manager,
-            int deviceId) throws IOException, InterruptedException {
+            int deviceId) throws InterruptedException, IOException {
         while (true) {
-            if (manager.getRoots(deviceId).length == 0) {
+            MtpDeviceRecord device = null;
+            for (final MtpDeviceRecord deviceCandidate : manager.getDevices()) {
+                if (deviceCandidate.deviceId == deviceId) {
+                    device = deviceCandidate;
+                    break;
+                }
+            }
+            if (device == null) {
+                throw new IOException("Device was detached.");
+            }
+            if (device.roots.length == 0) {
                 instrumentation.show("Wait for storages.");
                 Thread.sleep(1000);
                 continue;