Show MTP device as root when it has multiple/zero storages.

The CL updates MtpDocumentsProvider#queryRoots so that it fetches
devices from Database as well as storages when we don't have storages
under the device, or when we have multiple storages under the device.

BUG=26120019
Change-Id: Id2b140f00a1d49fa4da7e17d2564dbbaa1795e1e

7 files changed, 356 insertions, 223 deletions

diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
index cb49535ea55d..ac470671a582 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/Mapper.java
@@ -66,10 +66,13 @@ class Mapper {
         database.beginTransaction();
         try {
             final ContentValues[] valuesList = new ContentValues[1];
+            final ContentValues[] extraValuesList = new ContentValues[1];
             valuesList[0] = new ContentValues();
-            MtpDatabase.getDeviceDocumentValues(valuesList[0], device);
+            extraValuesList[0] = new ContentValues();
+            MtpDatabase.getDeviceDocumentValues(valuesList[0], extraValuesList[0], device);
             final boolean changed = putDocuments(
                     valuesList,
+                    extraValuesList,
                     COLUMN_PARENT_DOCUMENT_ID + " IS NULL",
                     EMPTY_ARGS,
                     /* heuristic */ false,
@@ -88,7 +91,7 @@ class Mapper {
      * @param roots List of root information.
      * @return If roots are added or removed from the database.
      */
-    synchronized boolean putRootDocuments(
+    synchronized boolean putStorageDocuments(
             String parentDocumentId, Resources resources, MtpRoot[] roots) {
         final SQLiteDatabase database = mDatabase.getSQLiteDatabase();
         database.beginTransaction();
@@ -109,36 +112,21 @@ class Mapper {
                     throw new Error("Unexpected map mode.");
             }
             final ContentValues[] valuesList = new ContentValues[roots.length];
+            final ContentValues[] extraValuesList = new ContentValues[roots.length];
             for (int i = 0; i < roots.length; i++) {
                 valuesList[i] = new ContentValues();
+                extraValuesList[i] = new ContentValues();
                 MtpDatabase.getStorageDocumentValues(
-                        valuesList[i], resources, parentDocumentId, roots[i]);
+                        valuesList[i], extraValuesList[i], resources, parentDocumentId, roots[i]);
             }
             final boolean changed = putDocuments(
                     valuesList,
+                    extraValuesList,
                     COLUMN_PARENT_DOCUMENT_ID + "=?",
                     strings(parentDocumentId),
                     heuristic,
                     mapColumn);
-            final ContentValues values = new ContentValues();
-            int i = 0;
-            for (final MtpRoot root : roots) {
-                // Use the same value for the root ID and the corresponding document ID.
-                final String documentId = valuesList[i++].getAsString(Document.COLUMN_DOCUMENT_ID);
-                // If it fails to insert/update documents, the document ID will be set with -1.
-                // In this case we don't insert/update root extra information neither.
-                if (documentId == null) {
-                    continue;
-                }
-                values.put(Root.COLUMN_ROOT_ID, documentId);
-                values.put(
-                        Root.COLUMN_FLAGS,
-                        Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE);
-                values.put(Root.COLUMN_AVAILABLE_BYTES, root.mFreeSpace);
-                values.put(Root.COLUMN_CAPACITY_BYTES, root.mMaxCapacity);
-                values.put(Root.COLUMN_MIME_TYPES, "");
-                database.replace(TABLE_ROOT_EXTRA, null, values);
-            }
+
             database.setTransactionSuccessful();
             return changed;
         } finally {
@@ -176,6 +164,7 @@ class Mapper {
         }
         putDocuments(
                 valuesList,
+                null,
                 COLUMN_PARENT_DOCUMENT_ID + "=?",
                 strings(parentId),
                 heuristic,
@@ -257,6 +246,7 @@ class Mapper {
      * 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.
+     * @param rootExtraValuesList Values for root extra to be stored in the database.
      * @param selection SQL where closure to select rows that shares the same parent.
      * @param args Argument for selection SQL.
      * @param heuristic Whether the mapping mode is heuristic.
@@ -264,6 +254,7 @@ class Mapper {
      */
     private boolean putDocuments(
             ContentValues[] valuesList,
+            @Nullable ContentValues[] rootExtraValuesList,
             String selection,
             String[] args,
             boolean heuristic,
@@ -272,7 +263,14 @@ class Mapper {
         boolean added = false;
         database.beginTransaction();
         try {
-            for (final ContentValues values : valuesList) {
+            for (int i = 0; i < valuesList.length; i++) {
+                final ContentValues values = valuesList[i];
+                final ContentValues rootExtraValues;
+                if (rootExtraValuesList != null) {
+                    rootExtraValues = rootExtraValuesList[i];
+                } else {
+                    rootExtraValues = null;
+                }
                 final Cursor candidateCursor = database.query(
                         TABLE_DOCUMENTS,
                         strings(Document.COLUMN_DOCUMENT_ID),
@@ -290,25 +288,26 @@ class Mapper {
                     final long rowId;
                     if (candidateCursor.getCount() == 0) {
                         rowId = database.insert(TABLE_DOCUMENTS, null, values);
-                        if (rowId == -1) {
-                            throw new SQLiteException("Failed to put a document into database.");
-                        }
                         added = true;
                     } else if (!heuristic) {
                         candidateCursor.moveToNext();
-                        final String documentId = candidateCursor.getString(0);
-                        rowId = database.update(
+                        rowId = candidateCursor.getLong(0);
+                        database.update(
                                 TABLE_DOCUMENTS,
                                 values,
                                 SELECTION_DOCUMENT_ID,
-                                strings(documentId));
+                                strings(rowId));
                     } else {
                         values.put(COLUMN_ROW_STATE, ROW_STATE_PENDING);
-                        rowId = database.insert(TABLE_DOCUMENTS, null, values);
+                        rowId = database.insertOrThrow(TABLE_DOCUMENTS, null, values);
                     }
                     // Document ID is a primary integer key of the table. So the returned row
                     // IDs should be same with the document ID.
                     values.put(Document.COLUMN_DOCUMENT_ID, rowId);
+                    if (rootExtraValues != null) {
+                        rootExtraValues.put(Root.COLUMN_ROOT_ID, rowId);
+                        database.replace(TABLE_ROOT_EXTRA, null, rootExtraValues);
+                    }
                 } finally {
                     candidateCursor.close();
                 }
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
index 10941eb02a0a..112914e93830 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
@@ -23,6 +23,9 @@ import android.content.ContentValues;
 import android.content.Context;
 import android.content.res.Resources;
 import android.database.Cursor;
+import android.database.DatabaseUtils;
+import android.database.MatrixCursor;
+import android.database.MatrixCursor.RowBuilder;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteOpenHelper;
 import android.database.sqlite.SQLiteQueryBuilder;
@@ -106,17 +109,93 @@ class MtpDatabase {
      * @return Database cursor.
      */
     Cursor queryRoots(String[] columnNames) {
-        final SQLiteQueryBuilder builder = new SQLiteQueryBuilder();
-        builder.setTables(JOIN_ROOTS);
-        builder.setProjectionMap(COLUMN_MAP_ROOTS);
-        return builder.query(
-                mDatabase,
-                columnNames,
-                COLUMN_ROW_STATE + " IN (?, ?) AND " + COLUMN_DOCUMENT_TYPE + " = ?",
-                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED, DOCUMENT_TYPE_STORAGE),
+        final String selection =
+                COLUMN_ROW_STATE + " IN (?, ?) AND " + COLUMN_DOCUMENT_TYPE + " = ?";
+        final Cursor deviceCursor = mDatabase.query(
+                TABLE_DOCUMENTS,
+                strings(COLUMN_DEVICE_ID),
+                selection,
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED, DOCUMENT_TYPE_DEVICE),
+                COLUMN_DEVICE_ID,
                 null,
                 null,
                 null);
+
+        try {
+            final SQLiteQueryBuilder builder = new SQLiteQueryBuilder();
+            builder.setTables(JOIN_ROOTS);
+            builder.setProjectionMap(COLUMN_MAP_ROOTS);
+            final MatrixCursor result = new MatrixCursor(columnNames);
+            final ContentValues values = new ContentValues();
+
+            while (deviceCursor.moveToNext()) {
+                final int deviceId = deviceCursor.getInt(0);
+                final Cursor storageCursor = builder.query(
+                        mDatabase,
+                        columnNames,
+                        selection + " AND " + COLUMN_DEVICE_ID + " = ?",
+                        strings(ROW_STATE_VALID,
+                                ROW_STATE_INVALIDATED,
+                                DOCUMENT_TYPE_STORAGE,
+                                deviceId),
+                        null,
+                        null,
+                        null);
+                try {
+                    values.clear();
+                    if (storageCursor.getCount() == 1) {
+                        storageCursor.moveToNext();
+                        DatabaseUtils.cursorRowToContentValues(storageCursor, values);
+                    } else {
+                        final Cursor cursor = builder.query(
+                                mDatabase,
+                                columnNames,
+                                selection + " AND " + COLUMN_DEVICE_ID + " = ?",
+                                strings(ROW_STATE_VALID,
+                                        ROW_STATE_INVALIDATED,
+                                        DOCUMENT_TYPE_DEVICE,
+                                        deviceId),
+                                null,
+                                null,
+                                null);
+                        try {
+                            cursor.moveToNext();
+                            DatabaseUtils.cursorRowToContentValues(cursor, values);
+                        } finally {
+                            cursor.close();
+                        }
+
+                        long capacityBytes = 0;
+                        long availableBytes = 0;
+                        int capacityIndex = cursor.getColumnIndex(Root.COLUMN_CAPACITY_BYTES);
+                        int availableIndex = cursor.getColumnIndex(Root.COLUMN_AVAILABLE_BYTES);
+                        while (storageCursor.moveToNext()) {
+                            // If requested columnNames does not include COLUMN_XXX_BYTES, we don't
+                            // calculate corresponding values.
+                            if (capacityIndex != -1) {
+                                capacityBytes += cursor.getLong(capacityIndex);
+                            }
+                            if (availableIndex != -1) {
+                                availableBytes += cursor.getLong(availableIndex);
+                            }
+                        }
+                        values.put(Root.COLUMN_CAPACITY_BYTES, capacityBytes);
+                        values.put(Root.COLUMN_AVAILABLE_BYTES, availableBytes);
+                    }
+                } finally {
+                    storageCursor.close();
+                }
+
+                final RowBuilder row = result.newRow();
+                for (final String key : values.keySet()) {
+                    row.add(key, values.get(key));
+                }
+            }
+
+            return result;
+        } finally {
+            deviceCursor.close();
+        }
     }
 
     /**
@@ -380,7 +459,10 @@ class MtpDatabase {
         context.deleteDatabase(DATABASE_NAME);
     }
 
-    static void getDeviceDocumentValues(ContentValues values, MtpDeviceRecord device) {
+    static void getDeviceDocumentValues(
+            ContentValues values,
+            ContentValues extraValues,
+            MtpDeviceRecord device) {
         values.clear();
         values.put(COLUMN_DEVICE_ID, device.deviceId);
         values.putNull(COLUMN_STORAGE_ID);
@@ -394,11 +476,15 @@ class MtpDatabase {
         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 : device.roots) {
-            size += root.mMaxCapacity - root.mFreeSpace;
-        }
-        values.put(Document.COLUMN_SIZE, size);
+        values.putNull(Document.COLUMN_SIZE);
+
+        extraValues.clear();
+        extraValues.put(
+                Root.COLUMN_FLAGS,
+                Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE);
+        extraValues.putNull(Root.COLUMN_AVAILABLE_BYTES);
+        extraValues.putNull(Root.COLUMN_CAPACITY_BYTES);
+        extraValues.put(Root.COLUMN_MIME_TYPES, "");
     }
 
     /**
@@ -408,7 +494,11 @@ class MtpDatabase {
      * @param root Root to be converted {@link ContentValues}.
      */
     static void getStorageDocumentValues(
-            ContentValues values, Resources resources, String parentDocumentId, MtpRoot root) {
+            ContentValues values,
+            ContentValues extraValues,
+            Resources resources,
+            String parentDocumentId,
+            MtpRoot root) {
         values.clear();
         values.put(COLUMN_DEVICE_ID, root.mDeviceId);
         values.put(COLUMN_STORAGE_ID, root.mStorageId);
@@ -424,6 +514,13 @@ class MtpDatabase {
         values.put(Document.COLUMN_FLAGS, 0);
         values.put(Document.COLUMN_SIZE,
                 (int) Math.min(root.mMaxCapacity - root.mFreeSpace, Integer.MAX_VALUE));
+
+        extraValues.put(
+                Root.COLUMN_FLAGS,
+                Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE);
+        extraValues.put(Root.COLUMN_AVAILABLE_BYTES, root.mFreeSpace);
+        extraValues.put(Root.COLUMN_CAPACITY_BYTES, root.mMaxCapacity);
+        extraValues.put(Root.COLUMN_MIME_TYPES, "");
     }
 
     /**
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
index a23358935987..f252e0f7d337 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabaseConstants.java
@@ -126,14 +126,14 @@ class MtpDatabaseConstants {
             Document.COLUMN_LAST_MODIFIED + " INTEGER," +
             Document.COLUMN_ICON + " INTEGER," +
             Document.COLUMN_FLAGS + " INTEGER NOT NULL," +
-            Document.COLUMN_SIZE + " INTEGER NOT NULL);";
+            Document.COLUMN_SIZE + " INTEGER);";
 
     static final String QUERY_CREATE_ROOT_EXTRA =
             "CREATE TABLE " + TABLE_ROOT_EXTRA + " (" +
             Root.COLUMN_ROOT_ID + " INTEGER PRIMARY KEY," +
             Root.COLUMN_FLAGS + " INTEGER NOT NULL," +
-            Root.COLUMN_AVAILABLE_BYTES + " INTEGER NOT NULL," +
-            Root.COLUMN_CAPACITY_BYTES + " INTEGER NOT NULL," +
+            Root.COLUMN_AVAILABLE_BYTES + " INTEGER," +
+            Root.COLUMN_CAPACITY_BYTES + " INTEGER," +
             Root.COLUMN_MIME_TYPES + " TEXT NOT NULL);";
 
     /**
@@ -145,18 +145,26 @@ class MtpDatabaseConstants {
         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);
+                Root.COLUMN_ICON,
+                TABLE_DOCUMENTS + "." + Document.COLUMN_ICON + " AS " + Root.COLUMN_ICON);
         COLUMN_MAP_ROOTS.put(
-                Root.COLUMN_DOCUMENT_ID, TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID);
+                Root.COLUMN_TITLE,
+                TABLE_DOCUMENTS + "." + Document.COLUMN_DISPLAY_NAME + " AS " + Root.COLUMN_TITLE);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_SUMMARY,
+                TABLE_DOCUMENTS + "." + Document.COLUMN_SUMMARY + " AS " + Root.COLUMN_SUMMARY);
+        COLUMN_MAP_ROOTS.put(
+                Root.COLUMN_DOCUMENT_ID,
+                TABLE_DOCUMENTS + "." + Document.COLUMN_DOCUMENT_ID +
+                " AS " + Root.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);
+        COLUMN_MAP_ROOTS.put(COLUMN_DEVICE_ID, COLUMN_DEVICE_ID);
     }
 
     private static String createJoinFromClosure(
diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
index e6c2726f531a..619ef54ed90e 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/RootScanner.java
@@ -127,7 +127,7 @@ final class RootScanner {
                         continue;
                     }
                     mDatabase.getMapper().startAddingDocuments(documentId);
-                    if (mDatabase.getMapper().putRootDocuments(
+                    if (mDatabase.getMapper().putStorageDocuments(
                             documentId, mResources, device.roots)) {
                         changed = true;
                     }
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
index 7bd9a1790739..f37a55c9568d 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/DocumentLoaderTest.java
@@ -41,7 +41,7 @@ public class DocumentLoaderTest extends AndroidTestCase {
     public void setUp() {
         mDatabase = new MtpDatabase(getContext(), MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", new TestResources(), new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", new TestResources(), new MtpRoot[] {
                 new MtpRoot(0, 0, "Device", "Storage", 1000, 1000, "")
         });
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
index b74517574c2e..1e1ea0a748a8 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
@@ -74,20 +74,24 @@ public class MtpDatabaseTest extends AndroidTestCase {
         return cursor.getString(cursor.getColumnIndex(columnName));
     }
 
-    public void testPutRootDocuments() throws Exception {
-        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,"")
+    public void testPutSingleStorageDocuments() throws Exception {
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device", true, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
+                new MtpRoot(0, 1, "Device", "Storage", 1000, 2000, "")
         });
+        mDatabase.getMapper().stopAddingDocuments("1");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(COLUMN_NAMES);
-            assertEquals(3, cursor.getCount());
+            assertEquals(1, cursor.getCount());
 
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(0, getInt(cursor, COLUMN_DEVICE_ID));
             assertEquals(1, getInt(cursor, COLUMN_STORAGE_ID));
             assertTrue(isNull(cursor, COLUMN_OBJECT_HANDLE));
@@ -101,14 +105,6 @@ public class MtpDatabaseTest extends AndroidTestCase {
             assertEquals(
                     MtpDatabaseConstants.DOCUMENT_TYPE_STORAGE, getInt(cursor, COLUMN_DOCUMENT_TYPE));
 
-            cursor.moveToNext();
-            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
-            assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
-
-            cursor.moveToNext();
-            assertEquals(3, getInt(cursor, COLUMN_DOCUMENT_ID));
-            assertEquals("Device /@#%&<>Storage", getString(cursor, COLUMN_DISPLAY_NAME));
-
             cursor.close();
         }
 
@@ -123,37 +119,58 @@ public class MtpDatabaseTest extends AndroidTestCase {
                     Root.COLUMN_AVAILABLE_BYTES,
                     Root.COLUMN_CAPACITY_BYTES
             });
+            assertEquals(1, cursor.getCount());
+
+            cursor.moveToNext();
+            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(
+                    Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE,
+                    getInt(cursor, Root.COLUMN_FLAGS));
+            assertEquals(R.drawable.ic_root_mtp, getInt(cursor, Root.COLUMN_ICON));
+            assertEquals("Device Storage", getString(cursor, Root.COLUMN_TITLE));
+            assertTrue(isNull(cursor, Root.COLUMN_SUMMARY));
+            assertEquals(2, getInt(cursor, Root.COLUMN_DOCUMENT_ID));
+            assertEquals(1000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
+            assertEquals(2000, getInt(cursor, Root.COLUMN_CAPACITY_BYTES));
+
+            cursor.close();
+        }
+    }
+
+    public void testPutStorageDocuments() throws Exception {
+        mDatabase.getMapper().startAddingDocuments("deviceDocId");
+        mDatabase.getMapper().putStorageDocuments("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,"")
+        });
+
+        {
+            final Cursor cursor = mDatabase.queryRootDocuments(COLUMN_NAMES);
             assertEquals(3, cursor.getCount());
 
             cursor.moveToNext();
-            assertEquals(1, cursor.getInt(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 Storage", cursor.getString(3));
-            assertTrue(cursor.isNull(4));
-            assertEquals(1, cursor.getInt(5));
-            assertEquals(1000, cursor.getInt(6));
-            assertEquals(2000, cursor.getInt(7));
+            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(0, getInt(cursor, COLUMN_DEVICE_ID));
+            assertEquals(1, getInt(cursor, COLUMN_STORAGE_ID));
+            assertTrue(isNull(cursor, COLUMN_OBJECT_HANDLE));
+            assertEquals(DocumentsContract.Document.MIME_TYPE_DIR, getString(cursor, COLUMN_MIME_TYPE));
+            assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
+            assertTrue(isNull(cursor, COLUMN_SUMMARY));
+            assertTrue(isNull(cursor, COLUMN_LAST_MODIFIED));
+            assertEquals(R.drawable.ic_root_mtp, getInt(cursor, COLUMN_ICON));
+            assertEquals(0, getInt(cursor, COLUMN_FLAGS));
+            assertEquals(1000, getInt(cursor, COLUMN_SIZE));
+            assertEquals(
+                    MtpDatabaseConstants.DOCUMENT_TYPE_STORAGE, getInt(cursor, COLUMN_DOCUMENT_TYPE));
 
             cursor.moveToNext();
-            assertEquals(2, cursor.getInt(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 Storage", cursor.getString(3));
-            assertTrue(cursor.isNull(4));
-            assertEquals(2, cursor.getInt(5));
-            assertEquals(2000, cursor.getInt(6));
-            assertEquals(4000, cursor.getInt(7));
+            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
 
             cursor.moveToNext();
-            assertEquals(3, cursor.getInt(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 /@#%&<>Storage", cursor.getString(3));
-            assertTrue(cursor.isNull(4));
-            assertEquals(3, cursor.getInt(5));
-            assertEquals(3000, cursor.getInt(6));
-            assertEquals(6000, cursor.getInt(7));
+            assertEquals(3, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals("Device /@#%&<>Storage", getString(cursor, COLUMN_DISPLAY_NAME));
 
             cursor.close();
         }
@@ -245,13 +262,9 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 MtpDatabaseConstants.COLUMN_STORAGE_ID,
                 DocumentsContract.Document.COLUMN_DISPLAY_NAME
         };
-        final String[] rootColumns = new String[] {
-                Root.COLUMN_ROOT_ID,
-                Root.COLUMN_AVAILABLE_BYTES
-        };
 
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 1000, 0, ""),
                 new MtpRoot(0, 101, "Device", "Storage B", 1001, 0, "")
         });
@@ -270,18 +283,6 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        {
-            final Cursor cursor = mDatabase.queryRoots(rootColumns);
-            assertEquals(2, cursor.getCount());
-            cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1001, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.close();
-        }
-
         mDatabase.getMapper().clearMapping();
 
         {
@@ -298,20 +299,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        {
-            final Cursor cursor = mDatabase.queryRoots(rootColumns);
-            assertEquals(2, cursor.getCount());
-            cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1001, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.close();
-        }
-
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage A", 2000, 0, ""),
                 new MtpRoot(0, 202, "Device", "Storage C", 2002, 0, "")
         });
@@ -334,21 +323,6 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        {
-            final Cursor cursor = mDatabase.queryRoots(rootColumns);
-            assertEquals(3, cursor.getCount());
-            cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(1001, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(4, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(2002, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.close();
-        }
-
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
 
         {
@@ -364,18 +338,6 @@ public class MtpDatabaseTest extends AndroidTestCase {
             assertEquals("Device Storage C", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
         }
-
-        {
-            final Cursor cursor = mDatabase.queryRoots(rootColumns);
-            assertEquals(2, cursor.getCount());
-            cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(2000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(4, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(2002, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.close();
-        }
     }
 
     public void testRestoreIdForChildDocuments() throws Exception {
@@ -461,26 +423,33 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
-        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().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device A", true, new MtpRoot[0]));
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(1, "Device B", true, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().startAddingDocuments("2");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device A", "Storage", 0, 0, "")
         });
-        mDatabase.getMapper().putRootDocuments("deviceDocIdB", resources, new MtpRoot[] {
-                new MtpRoot(1, 100, "Device", "Storage", 0, 0, "")
+        mDatabase.getMapper().putStorageDocuments("2", resources, new MtpRoot[] {
+                new MtpRoot(1, 100, "Device B", "Storage", 0, 0, "")
         });
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
             assertEquals(2, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(3, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(100, getInt(cursor, COLUMN_STORAGE_ID));
-            assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
+            assertEquals("Device A Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.moveToNext();
-            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(4, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(100, getInt(cursor, COLUMN_STORAGE_ID));
-            assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
+            assertEquals("Device B Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
         }
 
@@ -488,36 +457,36 @@ public class MtpDatabaseTest extends AndroidTestCase {
             final Cursor cursor = mDatabase.queryRoots(rootColumns);
             assertEquals(2, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(3, getInt(cursor, Root.COLUMN_ROOT_ID));
             assertEquals(0, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(4, getInt(cursor, Root.COLUMN_ROOT_ID));
             assertEquals(0, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.close();
         }
 
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments("deviceDocIdA");
-        mDatabase.getMapper().startAddingDocuments("deviceDocIdB");
-        mDatabase.getMapper().putRootDocuments("deviceDocIdA", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().startAddingDocuments("2");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, "")
         });
-        mDatabase.getMapper().putRootDocuments("deviceDocIdB", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("2", resources, new MtpRoot[] {
                 new MtpRoot(1, 300, "Device", "Storage", 3000, 0, "")
         });
-        mDatabase.getMapper().stopAddingDocuments("deviceDocIdA");
-        mDatabase.getMapper().stopAddingDocuments("deviceDocIdB");
+        mDatabase.getMapper().stopAddingDocuments("1");
+        mDatabase.getMapper().stopAddingDocuments("2");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
             assertEquals(2, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(5, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(200, getInt(cursor, COLUMN_STORAGE_ID));
             assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.moveToNext();
-            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(6, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(300, getInt(cursor, COLUMN_STORAGE_ID));
             assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
@@ -527,10 +496,10 @@ public class MtpDatabaseTest extends AndroidTestCase {
             final Cursor cursor = mDatabase.queryRoots(rootColumns);
             assertEquals(2, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(5, getInt(cursor, Root.COLUMN_ROOT_ID));
             assertEquals(2000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(6, getInt(cursor, Root.COLUMN_ROOT_ID));
             assertEquals(3000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.close();
         }
@@ -591,29 +560,34 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_AVAILABLE_BYTES
         };
 
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device",  false,  new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 300, "Device", "Storage", 3000, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
+        mDatabase.getMapper().stopAddingDocuments("1");
 
         {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(300, getInt(cursor, COLUMN_STORAGE_ID));
             assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
@@ -622,7 +596,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             final Cursor cursor = mDatabase.queryRoots(rootColumns);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
             assertEquals(3000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.close();
         }
@@ -634,19 +608,15 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 MtpDatabaseConstants.COLUMN_STORAGE_ID,
                 DocumentsContract.Document.COLUMN_DISPLAY_NAME
         };
-        final String[] rootColumns = new String[] {
-                Root.COLUMN_ROOT_ID,
-                Root.COLUMN_AVAILABLE_BYTES
-        };
 
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
 
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
                 new MtpRoot(0, 201, "Device", "Storage", 2001, 0, ""),
         });
@@ -665,33 +635,27 @@ public class MtpDatabaseTest extends AndroidTestCase {
             assertEquals("Device Storage", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
         }
-        {
-            final Cursor cursor = mDatabase.queryRoots(rootColumns);
-            assertEquals(2, cursor.getCount());
-            cursor.moveToNext();
-            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(2000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.moveToNext();
-            assertEquals(3, getInt(cursor, Root.COLUMN_ROOT_ID));
-            assertEquals(2001, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
-            cursor.close();
-        }
     }
 
     public void testReplaceExistingRoots() {
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device", true, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
         // The client code should be able to replace existing rows with new information.
         // Add one.
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
+        mDatabase.getMapper().stopAddingDocuments("1");
         // Replace it.
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage B", 1000, 1000, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
+        mDatabase.getMapper().stopAddingDocuments("1");
         {
             final String[] columns = new String[] {
                     DocumentsContract.Document.COLUMN_DOCUMENT_ID,
@@ -701,7 +665,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             final Cursor cursor = mDatabase.queryRootDocuments(columns);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, COLUMN_DOCUMENT_ID));
+            assertEquals(2, getInt(cursor, COLUMN_DOCUMENT_ID));
             assertEquals(100, getInt(cursor, COLUMN_STORAGE_ID));
             assertEquals("Device Storage B", getString(cursor, COLUMN_DISPLAY_NAME));
             cursor.close();
@@ -709,12 +673,14 @@ public class MtpDatabaseTest extends AndroidTestCase {
         {
             final String[] columns = new String[] {
                     Root.COLUMN_ROOT_ID,
+                    Root.COLUMN_TITLE,
                     Root.COLUMN_AVAILABLE_BYTES
             };
             final Cursor cursor = mDatabase.queryRoots(columns);
             assertEquals(1, cursor.getCount());
             cursor.moveToNext();
-            assertEquals(1, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals(2, getInt(cursor, Root.COLUMN_ROOT_ID));
+            assertEquals("Device Storage B", getString(cursor, Root.COLUMN_TITLE));
             assertEquals(1000, getInt(cursor, Root.COLUMN_AVAILABLE_BYTES));
             cursor.close();
         }
@@ -723,8 +689,13 @@ 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("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device", true, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().clearMapping();
@@ -732,18 +703,19 @@ public class MtpDatabaseTest extends AndroidTestCase {
         assertEquals(1, oldCursor.getCount());
 
         // Add one.
-        mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 101, "Device", "Storage B", 1000, 1000, ""),
         });
         // Add one more before resolving unmapped documents.
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
                 new MtpRoot(0, 102, "Device", "Storage B", 1000, 1000, ""),
         });
-        mDatabase.getMapper().stopAddingDocuments("deviceDocId");
+        mDatabase.getMapper().stopAddingDocuments("1");
 
         // Because the roots shares the same name, the roots should have new IDs.
-        final Cursor newCursor = mDatabase.queryRoots(strings(Root.COLUMN_ROOT_ID));
+        final Cursor newCursor = mDatabase.queryChildDocuments(
+                strings(Document.COLUMN_DOCUMENT_ID), "1");
         assertEquals(2, newCursor.getCount());
         oldCursor.moveToNext();
         newCursor.moveToNext();
@@ -755,9 +727,9 @@ public class MtpDatabaseTest extends AndroidTestCase {
         newCursor.close();
     }
 
-    public void testQueryDocument() {
+    public void testQueryDocuments() {
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
@@ -765,13 +737,61 @@ public class MtpDatabaseTest extends AndroidTestCase {
         final Cursor cursor = mDatabase.queryDocument("1", strings(Document.COLUMN_DISPLAY_NAME));
         assertEquals(1, cursor.getCount());
         cursor.moveToNext();
-        assertEquals("Device Storage A", cursor.getString(0));
+        assertEquals("Device Storage A", getString(cursor, Document.COLUMN_DISPLAY_NAME));
         cursor.close();
     }
 
+    public void testQueryRoots() {
+        // Add device document.
+        mDatabase.getMapper().startAddingDocuments(null);
+        mDatabase.getMapper().putDeviceDocument(
+                new MtpDeviceRecord(0, "Device", false, new MtpRoot[0]));
+        mDatabase.getMapper().stopAddingDocuments(null);
+
+        // It the device does not have storages, it shows a device root.
+        {
+            final Cursor cursor = mDatabase.queryRoots(strings(Root.COLUMN_TITLE));
+            assertEquals(1, cursor.getCount());
+            cursor.moveToNext();
+            assertEquals("Device", cursor.getString(0));
+            cursor.close();
+        }
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage A", 0, 0, "")
+        });
+        mDatabase.getMapper().stopAddingDocuments("1");
+
+        // It the device has single storage, it shows a storage root.
+        {
+            final Cursor cursor = mDatabase.queryRoots(strings(Root.COLUMN_TITLE));
+            assertEquals(1, cursor.getCount());
+            cursor.moveToNext();
+            assertEquals("Device Storage A", cursor.getString(0));
+            cursor.close();
+        }
+
+        mDatabase.getMapper().startAddingDocuments("1");
+        mDatabase.getMapper().putStorageDocuments("1", resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
+                new MtpRoot(0, 101, "Device", "Storage B", 0, 0, "")
+        });
+        mDatabase.getMapper().stopAddingDocuments("1");
+
+        // It the device has multiple storages, it shows a device root.
+        {
+            final Cursor cursor = mDatabase.queryRoots(strings(Root.COLUMN_TITLE));
+            assertEquals(1, cursor.getCount());
+            cursor.moveToNext();
+            assertEquals("Device", cursor.getString(0));
+            cursor.close();
+        }
+    }
+
     public void testGetParentId() throws FileNotFoundException {
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
@@ -790,7 +810,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
     public void testDeleteDocument() {
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
@@ -834,7 +854,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
     public void testPutNewDocument() {
         mDatabase.getMapper().startAddingDocuments("deviceDocId");
-        mDatabase.getMapper().putRootDocuments("deviceDocId", resources, new MtpRoot[] {
+        mDatabase.getMapper().putStorageDocuments("deviceDocId", resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
         });
         mDatabase.getMapper().stopAddingDocuments("deviceDocId");
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
index 884b1e2ccb1d..71c48971ca32 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDocumentsProviderTest.java
@@ -147,7 +147,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
             mProvider.openDevice(0);
             mResolver.waitForNotification(ROOTS_URI, 1);
             final Cursor cursor = mProvider.queryRoots(null);
-            assertEquals(1, cursor.getCount());
+            assertEquals(2, cursor.getCount());
             cursor.moveToNext();
             assertEquals("3", cursor.getString(0));
             assertEquals(Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE, cursor.getInt(1));
@@ -176,7 +176,7 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
     public void testQueryRoots_error() throws Exception {
         setupProvider(MtpDatabaseConstants.FLAG_DATABASE_IN_MEMORY);
         mMtpManager.addValidDevice(
-                new MtpDeviceRecord(0, "Device", false /* unopened */, new MtpRoot[0]));
+                new MtpDeviceRecord(0, "Device A", false /* unopened */, new MtpRoot[0]));
         mMtpManager.addValidDevice(new MtpDeviceRecord(
                 1,
                 "Device",
@@ -197,7 +197,16 @@ public class MtpDocumentsProviderTest extends AndroidTestCase {
             mResolver.waitForNotification(ROOTS_URI, 1);
 
             final Cursor cursor = mProvider.queryRoots(null);
-            assertEquals(1, cursor.getCount());
+            assertEquals(2, cursor.getCount());
+
+            cursor.moveToNext();
+            assertEquals("1", 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", cursor.getString(3));
+            assertEquals("1", cursor.getString(4));
+            assertEquals(0, cursor.getInt(5));
+
             cursor.moveToNext();
             assertEquals("3", cursor.getString(0));
             assertEquals(Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE, cursor.getInt(1));