Merge "Add device document to MtpDatabase."

9 files changed, 213 insertions, 138 deletions

diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
index 15ebe993f9c9..322246ac2ccb 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
@@ -42,6 +42,7 @@ import static com.android.mtp.MtpDatabase.strings;
  * Also see the comments of {@link MtpDatabase}.
  */
 class Mapper {
+    private static final String[] EMPTY_ARGS = new String[0];
     private final MtpDatabase mDatabase;
 
     /**
@@ -55,21 +56,43 @@ class Mapper {
         mDatabase = database;
     }
 
+    synchronized String putDeviceDocument(int deviceId, String name, MtpRoot[] roots) {
+        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(
+                    valuesList,
+                    COLUMN_PARENT_DOCUMENT_ID + " IS NULL",
+                    EMPTY_ARGS,
+                    /* heuristic */ false,
+                    COLUMN_DEVICE_ID);
+            database.setTransactionSuccessful();
+            return valuesList[0].getAsString(Document.COLUMN_DOCUMENT_ID);
+        } finally {
+            database.endTransaction();
+        }
+    }
+
     /**
      * Puts root information to database.
-     * @param deviceId Device ID
+     * @param parentDocumentId Document ID of device document.
      * @param resources Resources required to localize root name.
      * @param roots List of root information.
      * @return If roots are added or removed from the database.
      */
-    synchronized boolean putRootDocuments(int deviceId, Resources resources, MtpRoot[] roots) {
+    synchronized boolean putRootDocuments(
+            String parentDocumentId, Resources resources, MtpRoot[] roots) {
         final SQLiteDatabase database = mDatabase.getSQLiteDatabase();
         database.beginTransaction();
         try {
             final boolean heuristic;
             final String mapColumn;
-            Preconditions.checkState(mMappingMode.containsKey(/* no parent for root */ null));
-            switch (mMappingMode.get(/* no parent for root */ null)) {
+            Preconditions.checkState(mMappingMode.containsKey(parentDocumentId));
+            switch (mMappingMode.get(parentDocumentId)) {
                 case MAP_BY_MTP_IDENTIFIER:
                     heuristic = false;
                     mapColumn = COLUMN_STORAGE_ID;
@@ -83,16 +106,14 @@ class Mapper {
             }
             final ContentValues[] valuesList = new ContentValues[roots.length];
             for (int i = 0; i < roots.length; i++) {
-                if (roots[i].mDeviceId != deviceId) {
-                    throw new IllegalArgumentException();
-                }
                 valuesList[i] = new ContentValues();
-                MtpDatabase.getRootDocumentValues(valuesList[i], resources, roots[i]);
+                MtpDatabase.getStorageDocumentValues(
+                        valuesList[i], resources, parentDocumentId, roots[i]);
             }
             final boolean changed = putDocuments(
                     valuesList,
-                    COLUMN_PARENT_DOCUMENT_ID + " IS NULL",
-                    new String[0],
+                    COLUMN_PARENT_DOCUMENT_ID + "=?",
+                    strings(parentDocumentId),
                     heuristic,
                     mapColumn);
             final ContentValues values = new ContentValues();
@@ -146,7 +167,7 @@ class Mapper {
         final ContentValues[] valuesList = new ContentValues[documents.length];
         for (int i = 0; i < documents.length; i++) {
             valuesList[i] = new ContentValues();
-            MtpDatabase.getChildDocumentValues(
+            MtpDatabase.getObjectDocumentValues(
                     valuesList[i], deviceId, parentId, documents[i]);
         }
         putDocuments(
@@ -193,7 +214,7 @@ class Mapper {
             args = strings(parentDocumentId);
         } else {
             selection = COLUMN_PARENT_DOCUMENT_ID + " IS NULL";
-            args = new String[0];
+            args = EMPTY_ARGS;
         }
 
         final SQLiteDatabase database = mDatabase.getSQLiteDatabase();
@@ -312,7 +333,7 @@ class Mapper {
             args = strings(parentId);
         } else {
             selection = COLUMN_PARENT_DOCUMENT_ID + " IS NULL";
-            args = new String[0];
+            args = EMPTY_ARGS;
         }
         final String groupKey;
         switch (mMappingMode.get(parentId)) {
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
index eac9b98dd7ce..1ba3e314ff04 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
@@ -22,6 +22,7 @@ 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;
@@ -35,9 +36,9 @@ import android.provider.DocumentsContract.Root;
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.FileNotFoundException;
-import java.util.HashSet;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
-import java.util.Set;
 
 /**
  * Database for MTP objects.
@@ -107,11 +108,14 @@ class MtpDatabase {
      * @return Database cursor.
      */
     Cursor queryRoots(String[] columnNames) {
-        return mDatabase.query(
-                VIEW_ROOTS,
+        final SQLiteQueryBuilder builder = new SQLiteQueryBuilder();
+        builder.setTables(JOIN_ROOTS);
+        builder.setProjectionMap(COLUMN_MAP_ROOTS);
+        return builder.query(
+                mDatabase,
                 columnNames,
-                COLUMN_ROW_STATE + " IN (?, ?)",
-                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED),
+                COLUMN_ROW_STATE + " IN (?, ?) AND " + COLUMN_DOCUMENT_TYPE + " = ?",
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED, DOCUMENT_TYPE_STORAGE),
                 null,
                 null,
                 null);
@@ -128,8 +132,8 @@ class MtpDatabase {
         return mDatabase.query(
                 TABLE_DOCUMENTS,
                 columnNames,
-                COLUMN_ROW_STATE + " IN (?, ?)",
-                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED),
+                COLUMN_ROW_STATE + " IN (?, ?) AND " + COLUMN_DOCUMENT_TYPE + "=?",
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED, DOCUMENT_TYPE_STORAGE),
                 null,
                 null,
                 null);
@@ -216,7 +220,7 @@ class MtpDatabase {
      */
     String putNewDocument(int deviceId, String parentDocumentId, MtpObjectInfo info) {
         final ContentValues values = new ContentValues();
-        getChildDocumentValues(values, deviceId, parentDocumentId, info);
+        getObjectDocumentValues(values, deviceId, parentDocumentId, info);
         mDatabase.beginTransaction();
         try {
             final long id = mDatabase.insert(TABLE_DOCUMENTS, null, values);
@@ -344,7 +348,6 @@ class MtpDatabase {
         public void onCreate(SQLiteDatabase db) {
             db.execSQL(QUERY_CREATE_DOCUMENTS);
             db.execSQL(QUERY_CREATE_ROOT_EXTRA);
-            db.execSQL(QUERY_CREATE_VIEW_ROOTS);
         }
 
         @Override
@@ -358,18 +361,41 @@ class MtpDatabase {
         context.deleteDatabase(DATABASE_NAME);
     }
 
+    static void getDeviceDocumentValues(
+            ContentValues values, int deviceId, String name, MtpRoot[] roots) {
+        values.clear();
+        values.put(COLUMN_DEVICE_ID, 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.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) {
+            size += root.mMaxCapacity - root.mFreeSpace;
+        }
+        values.put(Document.COLUMN_SIZE, size);
+    }
+
     /**
      * Gets {@link ContentValues} for the given root.
      * @param values {@link ContentValues} that receives values.
      * @param resources Resources used to get localized root name.
      * @param root Root to be converted {@link ContentValues}.
      */
-    static void getRootDocumentValues(ContentValues values, Resources resources, MtpRoot root) {
+    static void getStorageDocumentValues(
+            ContentValues values, Resources resources, String parentDocumentId, MtpRoot root) {
         values.clear();
         values.put(COLUMN_DEVICE_ID, root.mDeviceId);
         values.put(COLUMN_STORAGE_ID, root.mStorageId);
         values.putNull(COLUMN_OBJECT_HANDLE);
-        values.putNull(COLUMN_PARENT_DOCUMENT_ID);
+        values.put(COLUMN_PARENT_DOCUMENT_ID, parentDocumentId);
         values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
         values.put(COLUMN_DOCUMENT_TYPE, DOCUMENT_TYPE_STORAGE);
         values.put(Document.COLUMN_MIME_TYPE, Document.MIME_TYPE_DIR);
@@ -389,7 +415,7 @@ class MtpDatabase {
      * @param parentId Parent document ID of the object.
      * @param info MTP object info.
      */
-    static void getChildDocumentValues(
+    static void getObjectDocumentValues(
             ContentValues values, int deviceId, String parentId, MtpObjectInfo info) {
         values.clear();
         final String mimeType = info.getFormat() == MtpConstants.FORMAT_ASSOCIATION ?
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
index 43dc8f524bdb..a23358935987 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
@@ -16,9 +16,13 @@
 
 package com.android.mtp;
 
+import android.database.sqlite.SQLiteQueryBuilder;
 import android.provider.DocumentsContract.Document;
 import android.provider.DocumentsContract.Root;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Class containing MtpDatabase constants.
  */
@@ -41,9 +45,13 @@ class MtpDatabaseConstants {
     static final String TABLE_ROOT_EXTRA = "RootExtra";
 
     /**
-     * View to join Documents and RootExtra tables to provide roots information.
+     * 'FROM' closure of joining TABLE_DOCUMENTS and TABLE_ROOT_EXTRA.
      */
-    static final String VIEW_ROOTS = "Roots";
+    static final String JOIN_ROOTS = createJoinFromClosure(
+            TABLE_DOCUMENTS,
+            TABLE_ROOT_EXTRA,
+            Document.COLUMN_DOCUMENT_ID,
+            Root.COLUMN_ROOT_ID);
 
     static final String COLUMN_DEVICE_ID = "device_id";
     static final String COLUMN_STORAGE_ID = "storage_id";
@@ -86,8 +94,6 @@ class MtpDatabaseConstants {
 
     /**
      * Document that represents a MTP device.
-     * Note we have "device" document only when the device has multiple storage volumes. Otherwise
-     * we regard the single "storage" document as root.
      */
     static final int DOCUMENT_TYPE_DEVICE = 0;
 
@@ -131,30 +137,31 @@ class MtpDatabaseConstants {
             Root.COLUMN_MIME_TYPES + " TEXT NOT NULL);";
 
     /**
-     * Creates a view to join Documents table and RootExtra table on their primary keys to
-     * provide DocumentContract.Root equivalent information.
+     * Map for columns names to provide DocumentContract.Root compatible columns.
+     * @see SQLiteQueryBuilder#setProjectionMap(Map)
      */
-    static final String QUERY_CREATE_VIEW_ROOTS =
-            "CREATE VIEW " + VIEW_ROOTS + " AS SELECT " +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID + " AS " +
-                            Root.COLUMN_ROOT_ID + "," +
-                    TABLE_ROOT_EXTRA + "." + Root.COLUMN_FLAGS + "," +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_ICON + " AS " +
-                            Root.COLUMN_ICON + "," +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_DISPLAY_NAME + " AS " +
-                            Root.COLUMN_TITLE + "," +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_SUMMARY + " AS " +
-                            Root.COLUMN_SUMMARY + "," +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID + " AS " +
-                    Root.COLUMN_DOCUMENT_ID + "," +
-                    TABLE_ROOT_EXTRA + "." + Root.COLUMN_AVAILABLE_BYTES + "," +
-                    TABLE_ROOT_EXTRA + "." + Root.COLUMN_CAPACITY_BYTES + "," +
-                    TABLE_ROOT_EXTRA + "." + Root.COLUMN_MIME_TYPES + "," +
-                    TABLE_DOCUMENTS + "." + COLUMN_ROW_STATE +
-            " FROM " + TABLE_DOCUMENTS + " INNER JOIN " + TABLE_ROOT_EXTRA +
-            " ON " +
-                    COLUMN_PARENT_DOCUMENT_ID + " IS NULL AND " +
-                    TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID +
-                    "=" +
-                    TABLE_ROOT_EXTRA + "." + Root.COLUMN_ROOT_ID;
+    static final Map<String, String> COLUMN_MAP_ROOTS;
+    static {
+        COLUMN_MAP_ROOTS = new HashMap<>();
+        COLUMN_MAP_ROOTS.put(Root.COLUMN_ROOT_ID, TABLE_ROOT_EXTRA + "." + Root.COLUMN_ROOT_ID);
+        COLUMN_MAP_ROOTS.put(Root.COLUMN_FLAGS, TABLE_ROOT_EXTRA + "." + Root.COLUMN_FLAGS);
+        COLUMN_MAP_ROOTS.put(Root.COLUMN_ICON, TABLE_DOCUMENTS + "." + Document.COLUMN_ICON);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_TITLE, TABLE_DOCUMENTS + "." + Document.COLUMN_DISPLAY_NAME);
+        COLUMN_MAP_ROOTS.put(Root.COLUMN_SUMMARY, TABLE_DOCUMENTS + "." + Document.COLUMN_SUMMARY);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_DOCUMENT_ID, TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_AVAILABLE_BYTES, TABLE_ROOT_EXTRA + "." + Root.COLUMN_AVAILABLE_BYTES);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_CAPACITY_BYTES, TABLE_ROOT_EXTRA + "." + Root.COLUMN_CAPACITY_BYTES);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_MIME_TYPES, TABLE_ROOT_EXTRA + "." + Root.COLUMN_MIME_TYPES);
+    }
+
+    private static String createJoinFromClosure(
+            String table1, String table2, String column1, String column2) {
+        return table1 + " INNER JOIN " + table2 +
+                " ON " + table1 + "." + column1 + " = " + table2 + "." + column2;
+    }
 }
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
index 9338c1b307d8..caa202402791 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDocumentsProvider.java
@@ -190,9 +190,6 @@ public class MtpDocumentsProvider extends DocumentsProvider {
             getDocumentLoader(parentIdentifier).clearTask(parentIdentifier);
             notifyChildDocumentsChange(parentIdentifier.mDocumentId);
         } catch (IOException error) {
-            for (final StackTraceElement element : error.getStackTrace()) {
-                Log.e("hirono", element.toString());
-            }
             throw new FileNotFoundException(error.getMessage());
         }
     }
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
index 52a751b235a4..197dec868c9a 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
@@ -7,8 +7,11 @@ 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;
+import java.util.Map;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.FutureTask;
@@ -106,26 +109,42 @@ final class RootScanner {
             int pollingCount = 0;
             while (!Thread.interrupted()) {
                 final int[] deviceIds = mManager.getOpenedDeviceIds();
-                if (deviceIds.length == 0) {
-                    return;
-                }
+                final Map<String, MtpRoot[]> rootsMap = new HashMap<>();
                 boolean changed = false;
+
+                // Update devices.
                 mDatabase.getMapper().startAddingDocuments(null /* parentDocumentId */);
-                for (int deviceId : deviceIds) {
+                for (final int deviceId : deviceIds) {
                     try {
                         final MtpRoot[] roots = mManager.getRoots(deviceId);
-                        if (mDatabase.getMapper().putRootDocuments(deviceId, mResources, roots)) {
+                        final String id = mDatabase.getMapper().putDeviceDocument(
+                                deviceId,
+                                mManager.getDeviceName(deviceId),
+                                roots);
+                        if (id != null) {
                             changed = true;
+                            rootsMap.put(id, roots);
                         }
-                    } catch (IOException | SQLiteException exception) {
+                    } 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());
                     }
                 }
-                if (mDatabase.getMapper().stopAddingDocuments(null /* parentDocumentId */)) {
-                    changed = true;
+                mDatabase.getMapper().stopAddingDocuments(null /* parentDocumentId */);
+
+                // Update roots.
+                for (final String documentId : rootsMap.keySet()) {
+                    mDatabase.getMapper().startAddingDocuments(documentId);
+                    if (mDatabase.getMapper().putRootDocuments(
+                            documentId, mResources, rootsMap.get(documentId))) {
+                        changed = true;
+                    }
+                    if (mDatabase.getMapper().stopAddingDocuments(documentId)) {
+                        changed = true;
+                    }
                 }
+
                 if (changed) {
                     notifyChange();
                 }
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
index 22d954d980a6..7bd9a1790739 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
@@ -40,11 +40,11 @@ public class DocumentLoaderTest extends AndroidTestCase {
     @Override
     public void setUp() {
         mDatabase = new MtpDatabase(getContext(), MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, new TestResources(), new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", new TestResources(), new MtpRoot[] {
                 new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "")
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
         mManager = new BlockableTestMtpManager(getContext());
         mResolver = new TestContentResolver();
         mLoader = new DocumentLoader(mManager, mResolver, mDatabase);
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
index eac5c6532a95..6d57c5b087f8 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
@@ -75,8 +75,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
     }
 
     public void testPutRootDocuments() throws Exception {
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 1, "Device", "Storage", 1000, 2000, ""),
                 new MtpRoot(0, 2, "Device", "Storage", 2000, 4000, ""),
                 new MtpRoot(0, 3, "Device", "/@#%&<>Storage", 3000, 6000,"")
@@ -250,8 +250,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_AVAILABLE_BYTES
         };
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 1000, 0, ""),
                 new MtpRoot(0, 101, "Device", "Storage B", 1001, 0, "")
         });
@@ -310,8 +310,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage A", 2000, 0, ""),
                 new MtpRoot(0, 202, "Device", "Storage C", 2002, 0, "")
         });
@@ -349,7 +349,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
@@ -451,10 +451,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
         }
     }
 
-    /**
-     * TODO: Enable this test after introducing device documents.
-     */
-    public void disabled_testRestoreIdForDifferentDevices() throws Exception {
+    public void testRestoreIdForDifferentDevices() throws Exception {
         final String[] columns = new String[] {
                 DocumentsContract.Document.COLUMN_DOCUMENT_ID,
                 MtpDatabaseConstants.COLUMN_STORAGE_ID,
@@ -464,11 +461,12 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocIdA");
+        mDatabase.getMapper().startAddingDocuments("deviceDocIdB");
+        mDatabase.getMapper().putRootDocuments("deviceDocIdA", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, "")
         });
-        mDatabase.getMapper().putRootDocuments(1, resources, new MtpRoot[] {
+        mDatabase.getMapper().putRootDocuments("deviceDocIdB", resources, new MtpRoot[] {
                 new MtpRoot(1, 100, "Device", "Storage", 0, 0, "")
         });
 
@@ -500,14 +498,16 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocIdA");
+        mDatabase.getMapper().startAddingDocuments("deviceDocIdB");
+        mDatabase.getMapper().putRootDocuments("deviceDocIdA", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, "")
         });
-        mDatabase.getMapper().putRootDocuments(1, resources, new MtpRoot[] {
+        mDatabase.getMapper().putRootDocuments("deviceDocIdB", resources, new MtpRoot[] {
                 new MtpRoot(1, 300, "Device", "Storage", 3000, 0, "")
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocIdA");
+        mDatabase.getMapper().stopAddingDocuments("deviceDocIdB");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
@@ -591,23 +591,23 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_AVAILABLE_BYTES
         };
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 300, "Device", "Storage", 3000, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
@@ -639,18 +639,18 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_AVAILABLE_BYTES
         };
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
                 new MtpRoot(0, 201, "Device", "Storage", 2001, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
@@ -681,17 +681,17 @@ public class MtpDatabaseTest extends AndroidTestCase {
     public void testReplaceExistingRoots() {
         // The client code should be able to replace existing rows with new information.
         // Add one.
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
         // Replace it.
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage B", 1000, 1000, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
         {
             final String[] columns = new String[] {
                     DocumentsContract.Document.COLUMN_DOCUMENT_ID,
@@ -723,8 +723,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
     public void testFailToReplaceExisitingUnmappedRoots() {
         // The client code should not be able to replace rows before resolving 'unmapped' rows.
         // Add one.
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
@@ -732,15 +732,15 @@ public class MtpDatabaseTest extends AndroidTestCase {
         assertEquals(1, oldCursor.getCount());
 
         // Add one.
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 101, "Device", "Storage B", 1000, 1000, ""),
         });
         // Add one more before resolving unmapped documents.
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 102, "Device", "Storage B", 1000, 1000, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         // Because the roots shares the same name, the roots should have new IDs.
         final Cursor newCursor = mDatabase.queryRoots(strings(Root.COLUMN_ROOT_ID));
@@ -756,11 +756,11 @@ public class MtpDatabaseTest extends AndroidTestCase {
     }
 
     public void testQueryDocument() {
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         final Cursor cursor = mDatabase.queryDocument("1", strings(Document.COLUMN_DISPLAY_NAME));
         assertEquals(1, cursor.getCount());
@@ -770,11 +770,11 @@ public class MtpDatabaseTest extends AndroidTestCase {
     }
 
     public void testGetParentId() throws FileNotFoundException {
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         mDatabase.getMapper().startAddingDocuments("1");
         mDatabase.getMapper().putChildDocuments(
@@ -789,11 +789,11 @@ public class MtpDatabaseTest extends AndroidTestCase {
     }
 
     public void testDeleteDocument() {
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         mDatabase.getMapper().startAddingDocuments("1");
         mDatabase.getMapper().putChildDocuments(
@@ -833,11 +833,11 @@ public class MtpDatabaseTest extends AndroidTestCase {
     }
 
     public void testPutNewDocument() {
-        mDatabase.getMapper().startAddingDocuments(null);
-        mDatabase.getMapper().putRootDocuments(0, resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments(null);
+        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         assertEquals(
                 "2",
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
index 597d51e1003c..b0e9722c74bb 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
@@ -138,11 +138,11 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
             final Cursor cursor = mProvider.queryRoots(null);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals("1", cursor.getString(0));
+            assertEquals("2", 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("1", cursor.getString(4));
+            assertEquals("2", cursor.getString(4));
             assertEquals(1024, cursor.getInt(5));
         }
 
@@ -153,11 +153,11 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
             assertEquals(2, cursor.getCount());
             cursor.moveToNext();
             cursor.moveToNext();
-            assertEquals("2", cursor.getString(0));
+            assertEquals("4", 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("4", cursor.getString(4));
             assertEquals(2048, cursor.getInt(5));
         }
     }
@@ -185,11 +185,11 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
             final Cursor cursor = mProvider.queryRoots(null);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals("1", cursor.getString(0));
+            assertEquals("2", 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("1", cursor.getString(4));
+            assertEquals("2", cursor.getString(4));
             assertEquals(2048, cursor.getInt(5));
         }
     }
@@ -216,12 +216,12 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
                                 .build()
                 });
 
-        final Cursor cursor = mProvider.queryDocument("2", null);
+        final Cursor cursor = mProvider.queryDocument("3", null);
         assertEquals(1, cursor.getCount());
 
         cursor.moveToNext();
 
-        assertEquals("2", cursor.getString(0));
+        assertEquals("3", cursor.getString(0));
         assertEquals("image/jpeg", cursor.getString(1));
         assertEquals("image.jpg", cursor.getString(2));
         assertEquals(1422716400000L, cursor.getLong(3));
@@ -255,11 +255,11 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
                                 .build()
                 });
 
-        final Cursor cursor = mProvider.queryDocument("2", null);
+        final Cursor cursor = mProvider.queryDocument("3", null);
         assertEquals(1, cursor.getCount());
 
         cursor.moveToNext();
-        assertEquals("2", cursor.getString(0));
+        assertEquals("3", cursor.getString(0));
         assertEquals(DocumentsContract.Document.MIME_TYPE_DIR, cursor.getString(1));
         assertEquals("directory", cursor.getString(2));
         assertEquals(1422716400000L, cursor.getLong(3));
@@ -287,11 +287,11 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
                         4096 /* total space */,
                         "" /* no volume identifier */)
         });
-        final Cursor cursor = mProvider.queryDocument("1", null);
+        final Cursor cursor = mProvider.queryDocument("2", null);
         assertEquals(1, cursor.getCount());
 
         cursor.moveToNext();
-        assertEquals("1", cursor.getString(0));
+        assertEquals("2", cursor.getString(0));
         assertEquals(DocumentsContract.Document.MIME_TYPE_DIR, cursor.getString(1));
         assertEquals("Device A Storage A", cursor.getString(2));
         assertTrue(cursor.isNull(3));
@@ -325,7 +325,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
         assertEquals(1, cursor.getCount());
 
         assertTrue(cursor.moveToNext());
-        assertEquals("2", cursor.getString(0));
+        assertEquals("3", cursor.getString(0));
         assertEquals("image/jpeg", cursor.getString(1));
         assertEquals("image.jpg", cursor.getString(2));
         assertEquals(0, cursor.getLong(3));
@@ -376,7 +376,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
                     .build()
         });
 
-        mProvider.deleteDocument("2");
+        mProvider.deleteDocument("3");
         assertEquals(1, mResolver.getChangeCount(
                 DocumentsContract.buildChildDocumentsUri(
                         MtpDocumentsProvider.AUTHORITY, "1")));
@@ -398,7 +398,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
                     .build()
         });
         try {
-            mProvider.deleteDocument("3");
+            mProvider.deleteDocument("4");
             fail();
         } catch (Throwable e) {
             assertTrue(e instanceof IOException);
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
index 3833799bacc4..ddc18a4a28e1 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/TestMtpManager.java
@@ -196,4 +196,9 @@ public class TestMtpManager extends MtpManager {
         }
         return result;
     }
+
+    @Override
+    String getDeviceName(int deviceId) throws IOException {
+        return "Device";
+    }
 }