Revamp Storage Settings header part

1. Add new object StorageEntry to encapsulate VolumeInfo and
   unsupported DiskInfo and missing VolumeRecord.
2. Replaces StorageSummaryDonutPreference with UsageProgressBarPreference.
3. Add storage select spinner.
4. Add a "Free up storage" preference to replace "Manage storage" button.

Bug: 174964885
Test: atest com.android.settings.deviceinfo.storage
      atest com.android.settings.deviceinfo
      manual
      Insert an USB drive, select the drive in StorageDashboardFragment
      and observe UI.
Change-Id: I83877f76869414de4fb2788b6b18fe507aa5cfcf
Merged-In: I83877f76869414de4fb2788b6b18fe507aa5cfcf
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 75884fa..3de4750 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -3333,6 +3333,10 @@
     <string name="storage_menu_manage">Manage storage</string>
     <!-- Storage setting. Keywords for Free up space. [CHAR LIMIT=NONE] -->
     <string name="keywords_storage_menu_free">clean, storage</string>
+    <!-- Storage setting. Title for storage free up option. [CHAR LIMIT=30] -->
+    <string name="storage_free_up_space_title">Free up space</string>
+    <!-- Storage setting. Summary for storage free up option. [CHAR LIMIT=NONE] -->
+    <string name="storage_free_up_space_summary">Go to Files app to manage and free up space</string>
 
     <!-- Storage setting.  Title for USB transfer settings [CHAR LIMIT=30]-->
     <string name="storage_title_usb">USB computer connection</string>
@@ -11455,6 +11459,10 @@
     <!-- Follows the percent of storage used by a storage volume. Exposed inside of a donut graph. [CHAR LIMIT=7]-->
     <string name="storage_percent_full">used</string>
 
+    <!-- Summary of a single storage volume used space. [CHAR LIMIT=24] -->
+    <string name="storage_usage_summary"><xliff:g id="number" example="128">%1$s</xliff:g> <xliff:g id="unit" example="KB">%2$s</xliff:g> used</string>
+    <!-- Summary of a single storage volume total space. [CHAR LIMIT=24] -->
+    <string name="storage_total_summary">Total <xliff:g id="number" example="128">%1$s</xliff:g> <xliff:g id="unit" example="KB">%2$s</xliff:g></string>
 
     <!-- Label for button allow user to remove the instant app from the device. -->
     <string name="clear_instant_app_data">Clear app</string>
diff --git a/res/xml/storage_dashboard_fragment.xml b/res/xml/storage_dashboard_fragment.xml
index bc58d7e..b49228e 100644
--- a/res/xml/storage_dashboard_fragment.xml
+++ b/res/xml/storage_dashboard_fragment.xml
@@ -19,11 +19,22 @@
     xmlns:settings="http://schemas.android.com/apk/res-auto"
     android:title="@string/storage_settings"
     android:orderingFromXml="false">
-    <com.android.settings.deviceinfo.storage.StorageSummaryDonutPreference
-        android:key="storage_summary"
-        android:order="0"
+    <com.android.settingslib.widget.SettingsSpinnerPreference
+        android:key="storage_spinner"
+        android:order="-2"
         settings:searchable="false"
-        settings:controller="com.android.settings.deviceinfo.storage.StorageSummaryDonutPreferenceController"/>
+        settings:controller="com.android.settings.deviceinfo.storage.StorageSelectionPreferenceController"/>
+    <com.android.settingslib.widget.UsageProgressBarPreference
+        android:key="storage_summary"
+        android:order="-1"
+        settings:searchable="false"
+        settings:controller="com.android.settings.deviceinfo.storage.StorageUsageProgressBarPreferenceController"/>
+    <Preference
+        android:key="free_up_space"
+        android:order="0"
+        android:title="@string/storage_free_up_space_title"
+        android:summary="@string/storage_free_up_space_summary"
+        settings:allowDividerAbove="true"/>
     <com.android.settings.widget.PrimarySwitchPreference
         android:fragment="com.android.settings.deletionhelper.AutomaticStorageManagerSettings"
         android:key="toggle_asm"
@@ -74,4 +85,4 @@
         android:key="pref_secondary_users"
         android:title="@string/storage_other_users"
         android:order="200" />
-</PreferenceScreen>
\ No newline at end of file
+</PreferenceScreen>
diff --git a/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuController.java b/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuController.java
index 00a79a0..549a02f 100644
--- a/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuController.java
+++ b/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuController.java
@@ -82,4 +82,8 @@
         }
         return false;
     }
+
+    public void setVolumeInfo(VolumeInfo volumeInfo) {
+        mVolumeInfo = volumeInfo;
+    }
 }
diff --git a/src/com/android/settings/deviceinfo/StorageDashboardFragment.java b/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
index 10c3a43..c9d2c6d 100644
--- a/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
+++ b/src/com/android/settings/deviceinfo/StorageDashboardFragment.java
@@ -17,39 +17,55 @@
 package com.android.settings.deviceinfo;
 
 import android.app.Activity;
+import android.app.Dialog;
 import android.app.settings.SettingsEnums;
 import android.app.usage.StorageStatsManager;
 import android.content.Context;
+import android.content.DialogInterface;
+import android.content.Intent;
 import android.graphics.drawable.Drawable;
 import android.os.Bundle;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.os.storage.DiskInfo;
+import android.os.storage.StorageEventListener;
 import android.os.storage.StorageManager;
 import android.os.storage.VolumeInfo;
+import android.os.storage.VolumeRecord;
 import android.provider.SearchIndexableResource;
+import android.text.TextUtils;
 import android.util.SparseArray;
 import android.view.View;
 
 import androidx.annotation.VisibleForTesting;
+import androidx.appcompat.app.AlertDialog;
+import androidx.fragment.app.Fragment;
 import androidx.loader.app.LoaderManager;
 import androidx.loader.content.Loader;
 import androidx.preference.Preference;
 
 import com.android.settings.R;
 import com.android.settings.Utils;
+import com.android.settings.core.SubSettingLauncher;
+import com.android.settings.core.instrumentation.InstrumentedDialogFragment;
 import com.android.settings.dashboard.DashboardFragment;
 import com.android.settings.dashboard.profileselector.ProfileSelectFragment;
 import com.android.settings.deviceinfo.storage.AutomaticStorageManagementSwitchPreferenceController;
 import com.android.settings.deviceinfo.storage.CachedStorageValuesHelper;
 import com.android.settings.deviceinfo.storage.SecondaryUserController;
 import com.android.settings.deviceinfo.storage.StorageAsyncLoader;
+import com.android.settings.deviceinfo.storage.StorageEntry;
 import com.android.settings.deviceinfo.storage.StorageItemPreferenceController;
+import com.android.settings.deviceinfo.storage.StorageSelectionPreferenceController;
+import com.android.settings.deviceinfo.storage.StorageUsageProgressBarPreferenceController;
 import com.android.settings.deviceinfo.storage.UserIconLoader;
 import com.android.settings.deviceinfo.storage.VolumeSizesLoader;
+import com.android.settings.overlay.FeatureFactory;
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settings.widget.EntityHeaderController;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.core.AbstractPreferenceController;
+import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 import com.android.settingslib.deviceinfo.PrivateStorageInfo;
 import com.android.settingslib.deviceinfo.StorageManagerVolumeProvider;
 import com.android.settingslib.search.SearchIndexable;
@@ -57,48 +73,184 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.stream.Collectors;
 
 @SearchIndexable
 public class StorageDashboardFragment extends DashboardFragment
         implements
-        LoaderManager.LoaderCallbacks<SparseArray<StorageAsyncLoader.AppsStorageResult>> {
+        LoaderManager.LoaderCallbacks<SparseArray<StorageAsyncLoader.AppsStorageResult>>,
+        Preference.OnPreferenceClickListener {
     private static final String TAG = "StorageDashboardFrag";
     private static final String SUMMARY_PREF_KEY = "storage_summary";
+    private static final String FREE_UP_SPACE_PREF_KEY = "free_up_space";
+    private static final String SELECTED_STORAGE_ENTRY_KEY = "selected_storage_entry_key";
     private static final int STORAGE_JOB_ID = 0;
     private static final int ICON_JOB_ID = 1;
     private static final int VOLUME_SIZE_JOB_ID = 2;
 
-    private VolumeInfo mVolume;
+    private StorageManager mStorageManager;
+    private final List<StorageEntry> mStorageEntries = new ArrayList<>();
+    private StorageEntry mSelectedStorageEntry;
     private PrivateStorageInfo mStorageInfo;
     private SparseArray<StorageAsyncLoader.AppsStorageResult> mAppsResult;
     private CachedStorageValuesHelper mCachedStorageValuesHelper;
 
     private StorageItemPreferenceController mPreferenceController;
     private PrivateVolumeOptionMenuController mOptionMenuController;
+    private StorageSelectionPreferenceController mStorageSelectionController;
+    private StorageUsageProgressBarPreferenceController mStorageUsageProgressBarController;
     private List<AbstractPreferenceController> mSecondaryUsers;
     private boolean mPersonalOnly;
+    private Preference mFreeUpSpacePreference;
+
+    private final StorageEventListener mStorageEventListener = new StorageEventListener() {
+        @Override
+        public void onVolumeStateChanged(VolumeInfo volumeInfo, int oldState, int newState) {
+            if (!isInteresting(volumeInfo)) {
+                return;
+            }
+
+            final StorageEntry storageEntry = new StorageEntry(getContext(), volumeInfo);
+            switch (volumeInfo.getState()) {
+                case VolumeInfo.STATE_MOUNTED:
+                case VolumeInfo.STATE_MOUNTED_READ_ONLY:
+                case VolumeInfo.STATE_UNMOUNTABLE:
+                    if (!mStorageEntries.contains(storageEntry)) {
+                        mStorageEntries.add(storageEntry);
+                        refreshUi();
+                    }
+                    break;
+                case VolumeInfo.STATE_REMOVED:
+                case VolumeInfo.STATE_UNMOUNTED:
+                case VolumeInfo.STATE_BAD_REMOVAL:
+                case VolumeInfo.STATE_EJECTING:
+                    if (mStorageEntries.remove(storageEntry)) {
+                        if (mSelectedStorageEntry.equals(storageEntry)) {
+                            mSelectedStorageEntry =
+                                    StorageEntry.getDefaultInternalStorageEntry(getContext());
+                        }
+                        refreshUi();
+                    }
+                    break;
+                default:
+                    // Do nothing.
+            }
+        }
+
+        @Override
+        public void onVolumeRecordChanged(VolumeRecord volumeRecord) {
+            final StorageEntry storageEntry = new StorageEntry(volumeRecord);
+            if (!mStorageEntries.contains(storageEntry)) {
+                mStorageEntries.add(storageEntry);
+                refreshUi();
+            }
+        }
+
+        @Override
+        public void onVolumeForgotten(String fsUuid) {
+            final StorageEntry storageEntry = new StorageEntry(
+                    new VolumeRecord(VolumeInfo.TYPE_PUBLIC, fsUuid));
+            if (mStorageEntries.remove(storageEntry)) {
+                if (mSelectedStorageEntry.equals(storageEntry)) {
+                    mSelectedStorageEntry =
+                            StorageEntry.getDefaultInternalStorageEntry(getContext());
+                }
+                refreshUi();
+            }
+        }
+
+        @Override
+        public void onDiskScanned(DiskInfo disk, int volumeCount) {
+            if (!isInteresting(disk)) {
+                return;
+            }
+            final StorageEntry storageEntry = new StorageEntry(disk);
+            if (!mStorageEntries.contains(storageEntry)) {
+                mStorageEntries.add(storageEntry);
+                refreshUi();
+            }
+        }
+
+        @Override
+        public void onDiskDestroyed(DiskInfo disk) {
+            final StorageEntry storageEntry = new StorageEntry(disk);
+            if (mStorageEntries.remove(storageEntry)) {
+                if (mSelectedStorageEntry.equals(storageEntry)) {
+                    mSelectedStorageEntry =
+                            StorageEntry.getDefaultInternalStorageEntry(getContext());
+                }
+                refreshUi();
+            }
+        }
+    };
+
+    private static boolean isInteresting(VolumeInfo volumeInfo) {
+        switch (volumeInfo.getType()) {
+            case VolumeInfo.TYPE_PRIVATE:
+            case VolumeInfo.TYPE_PUBLIC:
+            case VolumeInfo.TYPE_STUB:
+                return true;
+            default:
+                return false;
+        }
+    }
+
+    // Only interested in unsupported disk.
+    private static boolean isInteresting(DiskInfo disk) {
+        return disk.volumeCount == 0 && disk.size > 0;
+    }
+
+    private void refreshUi() {
+        mStorageSelectionController.setStorageEntries(mStorageEntries);
+        mStorageSelectionController.setSelectedStorageEntry(mSelectedStorageEntry);
+        mStorageUsageProgressBarController.setSelectedStorageEntry(mSelectedStorageEntry);
+
+        mOptionMenuController.setVolumeInfo(mSelectedStorageEntry.getVolumeInfo());
+
+        mPreferenceController.setVolume(mSelectedStorageEntry.getVolumeInfo());
+
+        if (mSelectedStorageEntry.isMounted()) {
+            getLoaderManager().restartLoader(STORAGE_JOB_ID, Bundle.EMPTY, this);
+            getLoaderManager()
+                 .restartLoader(VOLUME_SIZE_JOB_ID, Bundle.EMPTY, new VolumeSizeCallbacks());
+            getLoaderManager().restartLoader(ICON_JOB_ID, Bundle.EMPTY, new IconLoaderCallbacks());
+        } else {
+            mPreferenceController.clearStorageSizeDisplay();
+        }
+    }
 
     @Override
     public void onCreate(Bundle icicle) {
         super.onCreate(icicle);
 
-        // Initialize the storage sizes that we can quickly calc.
         final Activity activity = getActivity();
-        StorageManager sm = activity.getSystemService(StorageManager.class);
-        mVolume = Utils.maybeInitializeVolume(sm, getArguments());
+        mStorageManager = activity.getSystemService(StorageManager.class);
         mPersonalOnly = getArguments().getInt(ProfileSelectFragment.EXTRA_PROFILE)
                 == ProfileSelectFragment.ProfileType.PERSONAL;
-        if (mVolume == null) {
-            activity.finish();
-            return;
+
+        if (icicle == null) {
+            final VolumeInfo specifiedVolumeInfo =
+                    Utils.maybeInitializeVolume(mStorageManager, getArguments());
+            mSelectedStorageEntry = specifiedVolumeInfo == null
+                    ? StorageEntry.getDefaultInternalStorageEntry(getContext())
+                    : new StorageEntry(getContext(), specifiedVolumeInfo);
+        } else {
+            mSelectedStorageEntry = icicle.getParcelable(SELECTED_STORAGE_ENTRY_KEY);
         }
+
+        initializePreference();
         initializeOptionsMenu(activity);
+    }
+
+    private void initializePreference() {
         if (mPersonalOnly) {
             final Preference summary = getPreferenceScreen().findPreference(SUMMARY_PREF_KEY);
             if (summary != null) {
                 summary.setVisible(false);
             }
         }
+        mFreeUpSpacePreference = getPreferenceScreen().findPreference(FREE_UP_SPACE_PREF_KEY);
+        mFreeUpSpacePreference.setOnPreferenceClickListener(this);
     }
 
     @Override
@@ -106,12 +258,33 @@
         super.onAttach(context);
         use(AutomaticStorageManagementSwitchPreferenceController.class).setFragmentManager(
                 getFragmentManager());
+        mStorageSelectionController = use(StorageSelectionPreferenceController.class);
+        mStorageSelectionController.setOnItemSelectedListener(storageEntry -> {
+            mSelectedStorageEntry = storageEntry;
+            refreshUi();
+
+            if (storageEntry.isUnsupportedDiskInfo() || storageEntry.isUnmountable()) {
+                DiskInitFragment.show(this, R.string.storage_dialog_unmountable,
+                        storageEntry.getDiskId());
+            } else if (storageEntry.isMissingVolumeRecord()) {
+                final Bundle args = new Bundle();
+                args.putString(VolumeRecord.EXTRA_FS_UUID, storageEntry.getFsUuid());
+                new SubSettingLauncher(getContext())
+                        .setDestination(PrivateVolumeForget.class.getCanonicalName())
+                        .setTitleRes(R.string.storage_menu_forget)
+                        .setSourceMetricsCategory(getMetricsCategory())
+                        .setArguments(args)
+                        .launch();
+            }
+        });
+        mStorageUsageProgressBarController = use(StorageUsageProgressBarPreferenceController.class);
     }
 
     @VisibleForTesting
     void initializeOptionsMenu(Activity activity) {
-        mOptionMenuController = new PrivateVolumeOptionMenuController(
-                activity, mVolume, activity.getPackageManager());
+        mOptionMenuController = new PrivateVolumeOptionMenuController(activity,
+                mSelectedStorageEntry.getVolumeInfo(),
+                activity.getPackageManager());
         getSettingsLifecycle().addObserver(mOptionMenuController);
         setHasOptionsMenu(true);
         activity.invalidateOptionsMenu();
@@ -133,10 +306,37 @@
     @Override
     public void onResume() {
         super.onResume();
-        getLoaderManager().restartLoader(STORAGE_JOB_ID, Bundle.EMPTY, this);
-        getLoaderManager()
-                .restartLoader(VOLUME_SIZE_JOB_ID, Bundle.EMPTY, new VolumeSizeCallbacks());
-        getLoaderManager().restartLoader(ICON_JOB_ID, Bundle.EMPTY, new IconLoaderCallbacks());
+
+        mStorageEntries.clear();
+        mStorageEntries.addAll(mStorageManager.getVolumes().stream()
+                .filter(volumeInfo -> isInteresting(volumeInfo))
+                .map(volumeInfo -> new StorageEntry(getContext(), volumeInfo))
+                .collect(Collectors.toList()));
+        // Shows unsupported disks to give a chance to init.
+        mStorageEntries.addAll(mStorageManager.getDisks().stream()
+                .filter(disk -> isInteresting(disk))
+                .map(disk -> new StorageEntry(disk))
+                .collect(Collectors.toList()));
+        // Shows missing private volumes.
+        mStorageEntries.addAll(mStorageManager.getVolumeRecords().stream()
+                .filter(volumeRecord -> volumeRecord.getType() == VolumeInfo.TYPE_PRIVATE
+                        && mStorageManager.findVolumeByUuid(volumeRecord.getFsUuid()) == null)
+                .map(volumeRecord -> new StorageEntry(volumeRecord))
+                .collect(Collectors.toList()));
+        refreshUi();
+        mStorageManager.registerListener(mStorageEventListener);
+    }
+
+    @Override
+    public void onPause() {
+        super.onPause();
+        mStorageManager.unregisterListener(mStorageEventListener);
+    }
+
+    @Override
+    public void onSaveInstanceState(Bundle outState) {
+        outState.putParcelable(SELECTED_STORAGE_ENTRY_KEY, mSelectedStorageEntry);
+        super.onSaveInstanceState(outState);
     }
 
     @Override
@@ -148,7 +348,7 @@
         boolean stopLoading = false;
         if (mStorageInfo != null) {
             long privateUsedBytes = mStorageInfo.totalBytes - mStorageInfo.freeBytes;
-            mPreferenceController.setVolume(mVolume);
+            mPreferenceController.setVolume(mSelectedStorageEntry.getVolumeInfo());
             mPreferenceController.setUsedSize(privateUsedBytes);
             mPreferenceController.setTotalSize(mStorageInfo.totalBytes);
             for (int i = 0, size = mSecondaryUsers.size(); i < size; i++) {
@@ -197,7 +397,7 @@
 
         StorageManager sm = context.getSystemService(StorageManager.class);
         mPreferenceController = new StorageItemPreferenceController(context, this,
-                mVolume, new StorageManagerVolumeProvider(sm));
+                null /* volume */, new StorageManagerVolumeProvider(sm));
         controllers.add(mPreferenceController);
 
         final UserManager userManager = context.getSystemService(UserManager.class);
@@ -209,7 +409,7 @@
 
     @VisibleForTesting
     protected void setVolume(VolumeInfo info) {
-        mVolume = info;
+        mSelectedStorageEntry = new StorageEntry(getContext(), info);
     }
 
     /**
@@ -260,7 +460,7 @@
             Bundle args) {
         final Context context = getContext();
         return new StorageAsyncLoader(context, context.getSystemService(UserManager.class),
-                mVolume.fsUuid,
+                mSelectedStorageEntry.getFsUuid(),
                 new StorageStatsSource(context),
                 context.getPackageManager());
     }
@@ -277,6 +477,21 @@
     public void onLoaderReset(Loader<SparseArray<StorageAsyncLoader.AppsStorageResult>> loader) {
     }
 
+    @Override
+    public boolean onPreferenceClick(Preference preference) {
+        if (preference == mFreeUpSpacePreference) {
+            final Context context = getContext();
+            final MetricsFeatureProvider metricsFeatureProvider =
+                    FeatureFactory.getFactory(context).getMetricsFeatureProvider();
+            metricsFeatureProvider.logClickedPreference(preference, getMetricsCategory());
+            metricsFeatureProvider.action(context, SettingsEnums.STORAGE_FREE_UP_SPACE_NOW);
+            final Intent intent = new Intent(StorageManager.ACTION_MANAGE_STORAGE);
+            context.startActivity(intent);
+            return true;
+        }
+        return false;
+    }
+
     @VisibleForTesting
     public void setCachedStorageValuesHelper(CachedStorageValuesHelper helper) {
         mCachedStorageValuesHelper = helper;
@@ -340,8 +555,9 @@
     }
 
     private boolean isQuotaSupported() {
-        final StorageStatsManager stats = getActivity().getSystemService(StorageStatsManager.class);
-        return stats.isQuotaSupported(mVolume.fsUuid);
+        return mSelectedStorageEntry.isMounted()
+                && getActivity().getSystemService(StorageStatsManager.class)
+                        .isQuotaSupported(mSelectedStorageEntry.getFsUuid());
     }
 
     /**
@@ -378,11 +594,12 @@
             implements LoaderManager.LoaderCallbacks<PrivateStorageInfo> {
         @Override
         public Loader<PrivateStorageInfo> onCreateLoader(int id, Bundle args) {
-            Context context = getContext();
-            StorageManager sm = context.getSystemService(StorageManager.class);
-            StorageManagerVolumeProvider smvp = new StorageManagerVolumeProvider(sm);
+            final Context context = getContext();
+            final StorageManagerVolumeProvider smvp =
+                    new StorageManagerVolumeProvider(mStorageManager);
             final StorageStatsManager stats = context.getSystemService(StorageStatsManager.class);
-            return new VolumeSizesLoader(context, smvp, stats, mVolume);
+            return new VolumeSizesLoader(context, smvp, stats,
+                    mSelectedStorageEntry.getVolumeInfo());
         }
 
         @Override
@@ -402,4 +619,52 @@
             onReceivedSizes();
         }
     }
+
+    /** A dialog which guides users to initialize a specified unsupported disk. */
+    public static class DiskInitFragment extends InstrumentedDialogFragment {
+
+        private static final String TAG_DISK_INIT = "disk_init";
+
+        @Override
+        public int getMetricsCategory() {
+            return SettingsEnums.DIALOG_VOLUME_INIT;
+        }
+
+        /** Shows the dialo for the specified diskId from DiskInfo. */
+        public static void show(Fragment parent, int resId, String diskId) {
+            final Bundle args = new Bundle();
+            args.putInt(Intent.EXTRA_TEXT, resId);
+            args.putString(DiskInfo.EXTRA_DISK_ID, diskId);
+
+            final DiskInitFragment dialog = new DiskInitFragment();
+            dialog.setArguments(args);
+            dialog.setTargetFragment(parent, 0);
+            dialog.show(parent.getFragmentManager(), TAG_DISK_INIT);
+        }
+
+        @Override
+        public Dialog onCreateDialog(Bundle savedInstanceState) {
+            final Context context = getActivity();
+            final StorageManager storageManager = context.getSystemService(StorageManager.class);
+            final int resId = getArguments().getInt(Intent.EXTRA_TEXT);
+            final String diskId = getArguments().getString(DiskInfo.EXTRA_DISK_ID);
+            final DiskInfo disk = storageManager.findDiskById(diskId);
+
+            final AlertDialog.Builder builder = new AlertDialog.Builder(context);
+            builder.setMessage(TextUtils.expandTemplate(getText(resId), disk.getDescription()));
+
+            builder.setPositiveButton(R.string.storage_menu_set_up,
+                    new DialogInterface.OnClickListener() {
+                        @Override
+                        public void onClick(DialogInterface dialog, int which) {
+                            final Intent intent = new Intent(context, StorageWizardInit.class);
+                            intent.putExtra(DiskInfo.EXTRA_DISK_ID, diskId);
+                            startActivity(intent);
+                        }
+                    });
+            builder.setNegativeButton(R.string.cancel, null);
+
+            return builder.create();
+        }
+    }
 }
diff --git a/src/com/android/settings/deviceinfo/storage/StorageEntry.java b/src/com/android/settings/deviceinfo/storage/StorageEntry.java
new file mode 100644
index 0000000..f86fa79
--- /dev/null
+++ b/src/com/android/settings/deviceinfo/storage/StorageEntry.java
@@ -0,0 +1,283 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.os.Parcel;
+import android.os.Parcelable;
+import android.os.storage.DiskInfo;
+import android.os.storage.StorageManager;
+import android.os.storage.VolumeInfo;
+import android.os.storage.VolumeRecord;
+import android.text.TextUtils;
+
+import java.io.File;
+
+/**
+ * This object contains a {@link VolumeInfo} for a mountable storage or a {@link DiskInfo} for an
+ * unsupported disk which is not able to be mounted automatically.
+ */
+public class StorageEntry implements Comparable<StorageEntry>, Parcelable {
+
+    private final VolumeInfo mVolumeInfo;
+    private final DiskInfo mUnsupportedDiskInfo;
+    private final VolumeRecord mMissingVolumeRecord;
+
+    private final String mVolumeInfoDescription;
+
+    public StorageEntry(@NonNull Context context, @NonNull VolumeInfo volumeInfo) {
+        mVolumeInfo = volumeInfo;
+        mUnsupportedDiskInfo = null;
+        mMissingVolumeRecord = null;
+        mVolumeInfoDescription = context.getSystemService(StorageManager.class)
+                .getBestVolumeDescription(mVolumeInfo);
+    }
+
+    public StorageEntry(@NonNull DiskInfo diskInfo) {
+        mVolumeInfo = null;
+        mUnsupportedDiskInfo = diskInfo;
+        mMissingVolumeRecord = null;
+        mVolumeInfoDescription = null;
+    }
+
+    public StorageEntry(@NonNull VolumeRecord volumeRecord) {
+        mVolumeInfo = null;
+        mUnsupportedDiskInfo = null;
+        mMissingVolumeRecord = volumeRecord;
+        mVolumeInfoDescription = null;
+    }
+
+    private StorageEntry(Parcel in) {
+        mVolumeInfo = in.readParcelable(VolumeInfo.class.getClassLoader());
+        mUnsupportedDiskInfo = in.readParcelable(DiskInfo.class.getClassLoader());
+        mMissingVolumeRecord = in.readParcelable(VolumeRecord.class.getClassLoader());
+        mVolumeInfoDescription = in.readString();
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel out, int flags) {
+        out.writeParcelable(mVolumeInfo, 0 /* parcelableFlags */);
+        out.writeParcelable(mUnsupportedDiskInfo, 0 /* parcelableFlags */);
+        out.writeParcelable(mMissingVolumeRecord , 0 /* parcelableFlags */);
+        out.writeString(mVolumeInfoDescription);
+    }
+
+    public static final Parcelable.Creator<StorageEntry> CREATOR =
+            new Parcelable.Creator<StorageEntry>() {
+                public StorageEntry createFromParcel(Parcel in) {
+                    return new StorageEntry(in);
+                }
+
+                public StorageEntry[] newArray(int size) {
+                    return new StorageEntry[size];
+                }
+            };
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == this) {
+            return true;
+        }
+        if (!(o instanceof StorageEntry)) {
+            return false;
+        }
+
+        final StorageEntry StorageEntry = (StorageEntry) o;
+        if (isVolumeInfo()) {
+            return mVolumeInfo.equals(StorageEntry.mVolumeInfo);
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.equals(StorageEntry.mUnsupportedDiskInfo);
+        }
+        return mMissingVolumeRecord.equals(StorageEntry.mMissingVolumeRecord);
+    }
+
+    @Override
+    public int hashCode() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.hashCode();
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.hashCode();
+        }
+        return mMissingVolumeRecord.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.toString();
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.toString();
+        }
+        return mMissingVolumeRecord.toString();
+    }
+
+    @Override
+    public int compareTo(StorageEntry other) {
+        if (isDefaultInternalStorage() && !other.isDefaultInternalStorage()) {
+            return -1;
+        }
+        if (!isDefaultInternalStorage() && other.isDefaultInternalStorage()) {
+            return 1;
+        }
+
+        if (isVolumeInfo() && !other.isVolumeInfo()) {
+            return -1;
+        }
+        if (!isVolumeInfo() && other.isVolumeInfo()) {
+            return 1;
+        }
+
+        if (isPrivate() && !other.isPrivate()) {
+            return -1;
+        }
+        if (!isPrivate() && other.isPrivate()) {
+            return 1;
+        }
+
+        if (isMounted() && !other.isMounted()) {
+            return -1;
+        }
+        if (!isMounted() && other.isMounted()) {
+            return 1;
+        }
+
+        if (!isMissingVolumeRecord() && other.isMissingVolumeRecord()) {
+            return -1;
+        }
+        if (isMissingVolumeRecord() && !other.isMissingVolumeRecord()) {
+            return 1;
+        }
+
+        if (getDescription() == null) {
+            return 1;
+        }
+        if (other.getDescription() == null) {
+            return -1;
+        }
+        return getDescription().compareTo(other.getDescription());
+    }
+
+    /** Returns default internal storage. */
+    public static StorageEntry getDefaultInternalStorageEntry(Context context) {
+        return new StorageEntry(context, context.getSystemService(StorageManager.class)
+                .findVolumeById(VolumeInfo.ID_PRIVATE_INTERNAL));
+    }
+
+    /** If it's a VolumeInfo. */
+    public boolean isVolumeInfo() {
+        return mVolumeInfo != null;
+    }
+
+    /** If it's an unsupported DiskInfo. */
+    public boolean isUnsupportedDiskInfo() {
+        return mUnsupportedDiskInfo != null;
+    }
+
+    /** If it's a missing VolumeRecord. */
+    public boolean isMissingVolumeRecord() {
+        return mMissingVolumeRecord != null;
+    }
+
+    /** If it's a default internal storage. */
+    public boolean isDefaultInternalStorage() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.getType() == VolumeInfo.TYPE_PRIVATE
+                    && TextUtils.equals(mVolumeInfo.getId(), VolumeInfo.ID_PRIVATE_INTERNAL);
+        }
+        return false;
+    }
+
+    /** If it's a mounted storage. */
+    public boolean isMounted() {
+        return mVolumeInfo == null ? false : (mVolumeInfo.getState() == VolumeInfo.STATE_MOUNTED
+                || mVolumeInfo.getState() == VolumeInfo.STATE_MOUNTED_READ_ONLY);
+    }
+
+    /** If it's an unmountable storage. */
+    public boolean isUnmountable() {
+        return mVolumeInfo == null ? false : mVolumeInfo.getState() == VolumeInfo.STATE_UNMOUNTABLE;
+    }
+
+    /** If it's a private storage. */
+    public boolean isPrivate() {
+        return mVolumeInfo == null ? false : mVolumeInfo.getType() == VolumeInfo.TYPE_PRIVATE;
+    }
+
+    /** Returns description. */
+    public String getDescription() {
+        if (isVolumeInfo()) {
+            return mVolumeInfoDescription;
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.getDescription();
+        }
+        return mMissingVolumeRecord.getNickname();
+    }
+
+    /** Returns ID. */
+    public String getId() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.getId();
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.getId();
+        }
+        return mMissingVolumeRecord.getFsUuid();
+    }
+
+    /** Returns disk ID. */
+    public String getDiskId() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.getDiskId();
+        }
+        if (isUnsupportedDiskInfo()) {
+            return mUnsupportedDiskInfo.getId();
+        }
+        return null;
+    }
+
+    /** Returns fsUuid. */
+    public String getFsUuid() {
+        if (isVolumeInfo()) {
+            return mVolumeInfo.getFsUuid();
+        }
+        if (isUnsupportedDiskInfo()) {
+            return null;
+        }
+        return mMissingVolumeRecord.getFsUuid();
+    }
+
+    /** Returns root file if it's a VolumeInfo. */
+    public File getPath() {
+        return mVolumeInfo == null ? null : mVolumeInfo.getPath();
+    }
+
+    /** Returns VolumeInfo of the StorageEntry. */
+    public VolumeInfo getVolumeInfo() {
+        return mVolumeInfo;
+    }
+}
+
diff --git a/src/com/android/settings/deviceinfo/storage/StorageItemPreferenceController.java b/src/com/android/settings/deviceinfo/storage/StorageItemPreferenceController.java
index c2a0b62..dba72ba 100644
--- a/src/com/android/settings/deviceinfo/storage/StorageItemPreferenceController.java
+++ b/src/com/android/settings/deviceinfo/storage/StorageItemPreferenceController.java
@@ -158,6 +158,9 @@
                 intent = getAppsIntent();
                 break;
             case FILES_KEY:
+                if (mVolume == null) {
+                    break;
+                }
                 intent = getFilesIntent();
                 FeatureFactory.getFactory(mContext).getMetricsFeatureProvider().action(
                         mContext, SettingsEnums.STORAGE_FILES);
@@ -293,6 +296,17 @@
         mTotalSize = totalSizeBytes;
     }
 
+    /** Set storage size to 0 for each preference. */
+    public void clearStorageSizeDisplay() {
+        mPhotoPreference.setStorageSize(0L, 0L);
+        mAudioPreference.setStorageSize(0L, 0L);
+        mGamePreference.setStorageSize(0L, 0L);
+        mMoviesPreference.setStorageSize(0L, 0L);
+        mAppPreference.setStorageSize(0L, 0L);
+        mFilePreference.setStorageSize(0L, 0L);
+        mSystemPreference.setStorageSize(0L, 0L);
+    }
+
     /**
      * Returns a list of keys used by this preference controller.
      */
diff --git a/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceController.java b/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceController.java
new file mode 100644
index 0000000..03fddec
--- /dev/null
+++ b/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceController.java
@@ -0,0 +1,151 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import android.content.Context;
+import android.view.View;
+import android.view.ViewGroup;
+import android.widget.AdapterView;
+import android.widget.TextView;
+
+import androidx.annotation.VisibleForTesting;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settings.core.BasePreferenceController;
+import com.android.settingslib.widget.SettingsSpinnerPreference;
+import com.android.settingslib.widget.settingsspinner.SettingsSpinnerAdapter;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Shows a spinner for users to select a storage volume.
+ */
+public class StorageSelectionPreferenceController extends BasePreferenceController implements
+        AdapterView.OnItemSelectedListener {
+
+    @VisibleForTesting
+    SettingsSpinnerPreference mSpinnerPreference;
+    @VisibleForTesting
+    StorageAdapter mStorageAdapter;
+
+    private final List<StorageEntry> mStorageEntries = new ArrayList<>();
+
+    /** The interface for spinner selection callback. */
+    public interface OnItemSelectedListener {
+        /** Callbacked when the spinner selection is changed. */
+        void onItemSelected(StorageEntry storageEntry);
+    }
+    private OnItemSelectedListener mOnItemSelectedListener;
+
+    public StorageSelectionPreferenceController(Context context, String key) {
+        super(context, key);
+
+        mStorageAdapter = new StorageAdapter(context);
+    }
+
+    public void setOnItemSelectedListener(OnItemSelectedListener listener) {
+        mOnItemSelectedListener = listener;
+    }
+
+    /** Set the storages in the spinner. */
+    public void setStorageEntries(List<StorageEntry> storageEntries) {
+        mStorageAdapter.clear();
+        mStorageEntries.clear();
+        if (storageEntries == null || storageEntries.isEmpty()) {
+            return;
+        }
+        Collections.sort(mStorageEntries);
+        mStorageEntries.addAll(storageEntries);
+        mStorageAdapter.addAll(storageEntries);
+    }
+
+    /** set selected storage in the spinner. */
+    public void setSelectedStorageEntry(StorageEntry selectedStorageEntry) {
+        if (mSpinnerPreference == null || !mStorageEntries.contains(selectedStorageEntry)) {
+            return;
+        }
+        mSpinnerPreference.setSelection(mStorageAdapter.getPosition(selectedStorageEntry));
+    }
+
+    @Override
+    public int getAvailabilityStatus() {
+        return AVAILABLE_UNSEARCHABLE;
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        mSpinnerPreference = screen.findPreference(getPreferenceKey());
+        mSpinnerPreference.setAdapter(mStorageAdapter);
+        mSpinnerPreference.setOnItemSelectedListener(this);
+    }
+
+    @Override
+    public void onItemSelected(AdapterView<?> arg0, View arg1, int position, long id) {
+        if (mOnItemSelectedListener == null) {
+            return;
+        }
+        mOnItemSelectedListener.onItemSelected(mStorageAdapter.getItem(position));
+    }
+
+    @Override
+    public void onNothingSelected(AdapterView<?> arg0) {
+        // Do nothing.
+    }
+
+    @VisibleForTesting
+    class StorageAdapter extends SettingsSpinnerAdapter<StorageEntry> {
+
+        StorageAdapter(Context context) {
+            super(context);
+        }
+
+        @Override
+        public View getView(int position, View view, ViewGroup parent) {
+            if (view == null) {
+                view = getDefaultView(position, view, parent);
+            }
+
+            TextView textView = null;
+            try {
+                textView = (TextView) view;
+            } catch (ClassCastException e) {
+                throw new IllegalStateException("Default view should be a TextView, ", e);
+            }
+            textView.setText(getItem(position).getDescription());
+            return textView;
+        }
+
+        @Override
+        public View getDropDownView(int position, View view, ViewGroup parent) {
+            if (view == null) {
+                view = getDefaultDropDownView(position, view, parent);
+            }
+
+            TextView textView = null;
+            try {
+                textView = (TextView) view;
+            } catch (ClassCastException e) {
+                throw new IllegalStateException("Default drop down view should be a TextView, ", e);
+            }
+            textView.setText(getItem(position).getDescription());
+            return textView;
+        }
+    }
+}
+
diff --git a/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceController.java b/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceController.java
new file mode 100644
index 0000000..a00b25a
--- /dev/null
+++ b/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceController.java
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import android.app.usage.StorageStatsManager;
+import android.content.Context;
+import android.text.format.Formatter;
+import android.util.Log;
+
+import androidx.annotation.VisibleForTesting;
+import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settings.R;
+import com.android.settings.core.BasePreferenceController;
+import com.android.settingslib.utils.ThreadUtils;
+import com.android.settingslib.widget.UsageProgressBarPreference;
+
+import java.io.File;
+import java.io.IOException;
+
+/**
+ * Shows storage summary and progress.
+ */
+public class StorageUsageProgressBarPreferenceController extends BasePreferenceController {
+
+    private static final String TAG = "StorageProgressCtrl";
+
+    private final StorageStatsManager mStorageStatsManager;
+    @VisibleForTesting
+    long mUsedBytes;
+    @VisibleForTesting
+    long mTotalBytes;
+    private UsageProgressBarPreference mUsageProgressBarPreference;
+    private StorageEntry mStorageEntry;
+
+    public StorageUsageProgressBarPreferenceController(Context context, String key) {
+        super(context, key);
+
+        mStorageStatsManager = context.getSystemService(StorageStatsManager.class);
+    }
+
+    /** Set StorageEntry to display. */
+    public void setSelectedStorageEntry(StorageEntry storageEntry) {
+        mStorageEntry = storageEntry;
+        getStorageStatsAndUpdateUi();
+    }
+
+    @Override
+    public int getAvailabilityStatus() {
+        return AVAILABLE_UNSEARCHABLE;
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        mUsageProgressBarPreference = screen.findPreference(getPreferenceKey());
+        getStorageStatsAndUpdateUi();
+    }
+
+    private void getStorageStatsAndUpdateUi() {
+        ThreadUtils.postOnBackgroundThread(() -> {
+            try {
+                if (mStorageEntry == null || !mStorageEntry.isMounted()) {
+                    throw new IOException();
+                }
+
+                if (mStorageEntry.isPrivate()) {
+                    // StorageStatsManager can only query private storages.
+                    mTotalBytes = mStorageStatsManager.getTotalBytes(mStorageEntry.getFsUuid());
+                    mUsedBytes = mTotalBytes
+                            - mStorageStatsManager.getFreeBytes(mStorageEntry.getFsUuid());
+                } else {
+                    final File rootFile = mStorageEntry.getPath();
+                    if (rootFile == null) {
+                        Log.d(TAG, "Mounted public storage has null root path: " + mStorageEntry);
+                        throw new IOException();
+                    }
+                    mTotalBytes = rootFile.getTotalSpace();
+                    mUsedBytes = mTotalBytes - rootFile.getFreeSpace();
+                }
+            } catch (IOException e) {
+                // The storage device isn't present.
+                mTotalBytes = 0;
+                mUsedBytes = 0;
+            }
+
+            if (mUsageProgressBarPreference == null) {
+                return;
+            }
+            ThreadUtils.postOnMainThread(() ->
+                    updateState(mUsageProgressBarPreference)
+            );
+        });
+    }
+
+    @Override
+    public void updateState(Preference preference) {
+        mUsageProgressBarPreference.setUsageSummary(
+                getStorageSummary(R.string.storage_usage_summary, mUsedBytes));
+        mUsageProgressBarPreference.setTotalSummary(
+                getStorageSummary(R.string.storage_total_summary, mTotalBytes));
+        mUsageProgressBarPreference.setPercent(mUsedBytes, mTotalBytes);
+    }
+
+    private String getStorageSummary(int resId, long bytes) {
+        final Formatter.BytesResult result = Formatter.formatBytes(mContext.getResources(),
+                bytes, 0);
+        return mContext.getString(resId, result.value, result.units);
+    }
+}
diff --git a/src/com/android/settings/deviceinfo/storage/VolumeSizesLoader.java b/src/com/android/settings/deviceinfo/storage/VolumeSizesLoader.java
index d95befa..64510c6 100644
--- a/src/com/android/settings/deviceinfo/storage/VolumeSizesLoader.java
+++ b/src/com/android/settings/deviceinfo/storage/VolumeSizesLoader.java
@@ -26,6 +26,7 @@
 import com.android.settingslib.deviceinfo.StorageVolumeProvider;
 import com.android.settingslib.utils.AsyncLoaderCompat;
 
+import java.io.File;
 import java.io.IOException;
 
 public class VolumeSizesLoader extends AsyncLoaderCompat<PrivateStorageInfo> {
@@ -49,6 +50,11 @@
 
     @Override
     public PrivateStorageInfo loadInBackground() {
+        if (mVolume == null || (mVolume.getState() != VolumeInfo.STATE_MOUNTED
+                && mVolume.getState() != VolumeInfo.STATE_MOUNTED_READ_ONLY)) {
+            return new PrivateStorageInfo(0L /* freeBytes */, 0L /* totalBytes */);
+        }
+
         PrivateStorageInfo volumeSizes;
         try {
             volumeSizes = getVolumeSize(mVolumeProvider, mStats, mVolume);
@@ -62,8 +68,14 @@
     static PrivateStorageInfo getVolumeSize(
             StorageVolumeProvider storageVolumeProvider, StorageStatsManager stats, VolumeInfo info)
             throws IOException {
-        long privateTotalBytes = storageVolumeProvider.getTotalBytes(stats, info);
-        long privateFreeBytes = storageVolumeProvider.getFreeBytes(stats, info);
-        return new PrivateStorageInfo(privateFreeBytes, privateTotalBytes);
+        if (info.getType() == VolumeInfo.TYPE_PRIVATE) {
+            return new PrivateStorageInfo(storageVolumeProvider.getFreeBytes(stats, info),
+                    storageVolumeProvider.getTotalBytes(stats, info));
+        }
+        // TODO(b/174964885): It's confusing to use PrivateStorageInfo for a public storage,
+        //                    replace it with a new naming or a different object.
+        final File rootFile = info.getPath();
+        return rootFile == null ? new PrivateStorageInfo(0L /* freeBytes */, 0L /* totalBytes */)
+                : new PrivateStorageInfo(rootFile.getFreeSpace(), rootFile.getTotalSpace());
     }
 }
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index b55a788..b7ac4b1 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -32,6 +32,8 @@
         "platform-test-annotations",
         "truth-prebuilt",
         "ub-uiautomator",
+        "SettingsLibSettingsSpinner",
+        "SettingsLibUsageProgressBarPreference",
     ],
 
     // Include all test java files.
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageEntryTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageEntryTest.java
new file mode 100644
index 0000000..603d51e
--- /dev/null
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageEntryTest.java
@@ -0,0 +1,301 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.os.storage.DiskInfo;
+import android.os.storage.StorageManager;
+import android.os.storage.VolumeInfo;
+import android.os.storage.VolumeRecord;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.File;
+import java.util.Objects;
+
+@RunWith(AndroidJUnit4.class)
+public class StorageEntryTest {
+
+    private static final String VOLUME_INFO_ID = "volume_info_id";
+    private static final String DISK_INFO_ID = "disk_info_id";
+    private static final String VOLUME_RECORD_UUID = "volume_record_id";
+
+    @Mock
+    private VolumeInfo mVolumeInfo;
+    @Mock
+    private DiskInfo mDiskInfo;
+    @Mock
+    private VolumeRecord mVolumeRecord;
+
+    private Context mContext;
+    @Mock
+    private StorageManager mStorageManager;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        mContext = spy(ApplicationProvider.getApplicationContext());
+        when(mContext.getSystemService(StorageManager.class)).thenReturn(mStorageManager);
+    }
+
+    @Test
+    public void equals_volumesOfSameId_shouldBeTheSame() {
+        final StorageEntry volumeStorage1 = new StorageEntry(mContext,
+                new VolumeInfo(VOLUME_INFO_ID, 0 /* type */, null /* disk */, null /* partGuid */));
+        final StorageEntry volumeStorage2 = new StorageEntry(mContext,
+                new VolumeInfo(VOLUME_INFO_ID, 0 /* type */, null /* disk */, null /* partGuid */));
+        final StorageEntry diskStorage1 =
+                new StorageEntry(new DiskInfo(DISK_INFO_ID, 0 /* flags */));
+        final StorageEntry diskStorage2 =
+                new StorageEntry(new DiskInfo(DISK_INFO_ID, 0 /* flags */));
+        final StorageEntry volumeRecordStorage1 = new StorageEntry(new VolumeRecord(0 /* flags */,
+                VOLUME_RECORD_UUID));
+        final StorageEntry volumeRecordStorage2 = new StorageEntry(new VolumeRecord(0 /* flags */,
+                VOLUME_RECORD_UUID));
+
+        assertThat(Objects.equals(volumeStorage1, volumeStorage2)).isTrue();
+        assertThat(Objects.equals(diskStorage1, diskStorage2)).isTrue();
+        assertThat(Objects.equals(volumeRecordStorage1, volumeRecordStorage2)).isTrue();
+    }
+
+    @Test
+    public void equals_volumesOfDifferentId_shouldBeDifferent() {
+        final StorageEntry volumeStorage1 = new StorageEntry(mContext,
+                new VolumeInfo(VOLUME_INFO_ID, 0 /* type */, null /* disk */, null /* partGuid */));
+        final StorageEntry volumeStorage2 = new StorageEntry(mContext,
+                new VolumeInfo("id2", 0 /* type */, null /* disk */, null /* partGuid */));
+        final StorageEntry diskStorage1 =
+                new StorageEntry(new DiskInfo(DISK_INFO_ID, 0 /* flags */));
+        final StorageEntry diskStorage2 =
+                new StorageEntry(new DiskInfo("id2", 0 /* flags */));
+        final StorageEntry volumeRecordStorage1 = new StorageEntry(new VolumeRecord(0 /* flags */,
+                VOLUME_RECORD_UUID));
+        final StorageEntry volumeRecordStorage2 = new StorageEntry(new VolumeRecord(0 /* flags */,
+                "id2"));
+
+        assertThat(Objects.equals(volumeStorage1, volumeStorage2)).isFalse();
+        assertThat(Objects.equals(diskStorage1, diskStorage2)).isFalse();
+        assertThat(Objects.equals(volumeRecordStorage1, volumeRecordStorage2)).isFalse();
+    }
+
+    @Test
+    public void compareTo_defaultInternalStorage_shouldBeAtTopMost() {
+        final StorageEntry storage1 = mock(StorageEntry.class);
+        when(storage1.isDefaultInternalStorage()).thenReturn(true);
+        final StorageEntry storage2 = mock(StorageEntry.class);
+        when(storage2.isDefaultInternalStorage()).thenReturn(false);
+
+        assertThat(storage1.compareTo(storage2) > 0).isTrue();
+    }
+
+    @Test
+    public void getDefaultInternalStorageEntry_shouldReturnVolumeInfoStorageOfIdPrivateInternal() {
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        when(mStorageManager.findVolumeById(VolumeInfo.ID_PRIVATE_INTERNAL)).thenReturn(volumeInfo);
+
+        assertThat(StorageEntry.getDefaultInternalStorageEntry(mContext))
+                .isEqualTo(new StorageEntry(mContext, volumeInfo));
+    }
+
+    @Test
+    public void isVolumeInfo_shouldReturnTrueForVolumeInfo() {
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        final StorageEntry storage = new StorageEntry(mContext, volumeInfo);
+
+        assertThat(storage.isVolumeInfo()).isTrue();
+        assertThat(storage.isUnsupportedDiskInfo()).isFalse();
+        assertThat(storage.isMissingVolumeRecord()).isFalse();
+    }
+
+    @Test
+    public void isUnsupportedDiskInfo_shouldReturnTrueForDiskInfo() {
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry storage = new StorageEntry(diskInfo);
+
+        assertThat(storage.isVolumeInfo()).isFalse();
+        assertThat(storage.isUnsupportedDiskInfo()).isTrue();
+        assertThat(storage.isMissingVolumeRecord()).isFalse();
+    }
+
+    @Test
+    public void isMissingVolumeRecord_shouldReturnTrueForVolumeRecord() {
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry storage = new StorageEntry(volumeRecord);
+
+        assertThat(storage.isVolumeInfo()).isFalse();
+        assertThat(storage.isUnsupportedDiskInfo()).isFalse();
+        assertThat(storage.isMissingVolumeRecord()).isTrue();
+    }
+
+    @Test
+    public void isMounted_mountedOrMountedReadOnly_shouldReturnTrue() {
+        final VolumeInfo mountedVolumeInfo1 = mock(VolumeInfo.class);
+        final StorageEntry mountedStorage1 = new StorageEntry(mContext, mountedVolumeInfo1);
+        when(mountedVolumeInfo1.getState()).thenReturn(VolumeInfo.STATE_MOUNTED);
+        final VolumeInfo mountedVolumeInfo2 = mock(VolumeInfo.class);
+        when(mountedVolumeInfo2.getState()).thenReturn(VolumeInfo.STATE_MOUNTED_READ_ONLY);
+        final StorageEntry mountedStorage2 = new StorageEntry(mContext, mountedVolumeInfo2);
+
+        assertThat(mountedStorage1.isMounted()).isTrue();
+        assertThat(mountedStorage2.isMounted()).isTrue();
+    }
+
+    @Test
+    public void isMounted_nonVolumeInfo_shouldReturnFalse() {
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage2 = new StorageEntry(volumeRecord);
+
+        assertThat(diskStorage.isMounted()).isFalse();
+        assertThat(recordStorage2.isMounted()).isFalse();
+    }
+
+    @Test
+    public void isUnmountable_unmountableVolume_shouldReturnTrue() {
+        final VolumeInfo unmountableVolumeInfo = mock(VolumeInfo.class);
+        final StorageEntry mountedStorage = new StorageEntry(mContext, unmountableVolumeInfo);
+        when(unmountableVolumeInfo.getState()).thenReturn(VolumeInfo.STATE_UNMOUNTABLE);
+
+        assertThat(mountedStorage.isUnmountable()).isTrue();
+    }
+
+    @Test
+    public void isUnmountable_nonVolumeInfo_shouldReturnFalse() {
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage2 = new StorageEntry(volumeRecord);
+
+        assertThat(diskStorage.isUnmountable()).isFalse();
+        assertThat(recordStorage2.isUnmountable()).isFalse();
+    }
+
+    @Test
+    public void isPrivate_privateVolume_shouldReturnTrue() {
+        final VolumeInfo privateVolumeInfo = mock(VolumeInfo.class);
+        final StorageEntry privateStorage = new StorageEntry(mContext, privateVolumeInfo);
+        when(privateVolumeInfo.getType()).thenReturn(VolumeInfo.TYPE_PRIVATE);
+
+        assertThat(privateStorage.isPrivate()).isTrue();
+    }
+
+    @Test
+    public void isPrivate_nonVolumeInfo_shouldReturnFalse() {
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage2 = new StorageEntry(volumeRecord);
+
+        assertThat(diskStorage.isPrivate()).isFalse();
+        assertThat(recordStorage2.isPrivate()).isFalse();
+    }
+
+    @Test
+    public void getDescription_shouldReturnDescription() {
+        final String description = "description";
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        when(mStorageManager.getBestVolumeDescription(volumeInfo)).thenReturn(description);
+        final StorageEntry volumeStorage = new StorageEntry(mContext, volumeInfo);
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        when(diskInfo.getDescription()).thenReturn(description);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage = new StorageEntry(volumeRecord);
+        when(volumeRecord.getNickname()).thenReturn(description);
+
+        assertThat(volumeStorage.getDescription()).isEqualTo(description);
+        assertThat(diskStorage.getDescription()).isEqualTo(description);
+        assertThat(recordStorage.getDescription()).isEqualTo(description);
+    }
+
+    @Test
+    public void getDiskId_shouldReturnDiskId() {
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        final StorageEntry volumeStorage = new StorageEntry(mContext, volumeInfo);
+        when(volumeInfo.getDiskId()).thenReturn(VOLUME_INFO_ID);
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        when(diskInfo.getId()).thenReturn(DISK_INFO_ID);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage = new StorageEntry(volumeRecord);
+
+        assertThat(volumeStorage.getDiskId()).isEqualTo(VOLUME_INFO_ID);
+        assertThat(diskStorage.getDiskId()).isEqualTo(DISK_INFO_ID);
+        assertThat(recordStorage.getDiskId()).isEqualTo(null);
+    }
+
+    @Test
+    public void getFsUuid_shouldReturnFsUuid() {
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        final StorageEntry volumeStorage = new StorageEntry(mContext, volumeInfo);
+        when(volumeInfo.getFsUuid()).thenReturn(VOLUME_INFO_ID);
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage = new StorageEntry(volumeRecord);
+        when(volumeRecord.getFsUuid()).thenReturn(VOLUME_RECORD_UUID);
+
+        assertThat(volumeStorage.getFsUuid()).isEqualTo(VOLUME_INFO_ID);
+        assertThat(diskStorage.getFsUuid()).isEqualTo(null);
+        assertThat(recordStorage.getFsUuid()).isEqualTo(VOLUME_RECORD_UUID);
+    }
+
+    @Test
+    public void getPath_shouldReturnPath() {
+        final File file = new File("fakePath");
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        final StorageEntry volumeStorage = new StorageEntry(mContext, volumeInfo);
+        when(volumeInfo.getPath()).thenReturn(file);
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage = new StorageEntry(volumeRecord);
+
+        assertThat(volumeStorage.getPath()).isEqualTo(file);
+        assertThat(diskStorage.getPath()).isEqualTo(null);
+        assertThat(recordStorage.getPath()).isEqualTo(null);
+    }
+
+    @Test
+    public void getVolumeInfo_shouldVolumeInfo() {
+        final VolumeInfo volumeInfo = mock(VolumeInfo.class);
+        final StorageEntry volumeStorage = new StorageEntry(mContext, volumeInfo);
+        final DiskInfo diskInfo = mock(DiskInfo.class);
+        final StorageEntry diskStorage = new StorageEntry(diskInfo);
+        final VolumeRecord volumeRecord = mock(VolumeRecord.class);
+        final StorageEntry recordStorage = new StorageEntry(volumeRecord);
+
+        assertThat(volumeStorage.getVolumeInfo()).isEqualTo(volumeInfo);
+        assertThat(diskStorage.getVolumeInfo()).isEqualTo(null);
+        assertThat(recordStorage.getVolumeInfo()).isEqualTo(null);
+    }
+}
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceControllerTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceControllerTest.java
new file mode 100644
index 0000000..86351cb
--- /dev/null
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageSelectionPreferenceControllerTest.java
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.os.Looper;
+import android.os.storage.StorageManager;
+
+import androidx.preference.PreferenceManager;
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.settingslib.widget.SettingsSpinnerPreference;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+@RunWith(AndroidJUnit4.class)
+public class StorageSelectionPreferenceControllerTest {
+
+    private static final String PREFERENCE_KEY = "preference_key";
+
+    private Context mContext;
+    private StorageManager mStorageManager;
+    private StorageSelectionPreferenceController mController;
+
+    @Before
+    public void setUp() throws Exception {
+        mContext = ApplicationProvider.getApplicationContext();
+        mStorageManager = mContext.getSystemService(StorageManager.class);
+        mController = new StorageSelectionPreferenceController(mContext, PREFERENCE_KEY);
+    }
+
+    @Test
+    public void setStorageEntries_fromStorageManager_correctAdapterItems() {
+        final List<StorageEntry> storageEntries = mStorageManager.getVolumes().stream()
+                .map(volumeInfo -> new StorageEntry(mContext, volumeInfo))
+                .collect(Collectors.toList());
+
+        mController.setStorageEntries(storageEntries);
+
+        final int adapterItemCount = mController.mStorageAdapter.getCount();
+        assertThat(adapterItemCount).isEqualTo(storageEntries.size());
+        for (int i = 0; i < adapterItemCount; i++) {
+            assertThat(storageEntries.get(i).getDescription())
+                    .isEqualTo(mController.mStorageAdapter.getItem(i).getDescription());
+        }
+    }
+
+    @Test
+    public void setSelectedStorageEntry_primaryStorage_correctSelectedAdapterItem() {
+        if (Looper.myLooper() == null) {
+            Looper.prepare();
+        }
+        final PreferenceManager preferenceManager = new PreferenceManager(mContext);
+        final PreferenceScreen preferenceScreen =
+                preferenceManager.createPreferenceScreen(mContext);
+        final SettingsSpinnerPreference spinnerPreference = new SettingsSpinnerPreference(mContext);
+        spinnerPreference.setKey(PREFERENCE_KEY);
+        preferenceScreen.addPreference(spinnerPreference);
+        mController.displayPreference(preferenceScreen);
+        final StorageEntry primaryStorageEntry =
+                StorageEntry.getDefaultInternalStorageEntry(mContext);
+        mController.setStorageEntries(mStorageManager.getVolumes().stream()
+                .map(volumeInfo -> new StorageEntry(mContext, volumeInfo))
+                .collect(Collectors.toList()));
+
+        mController.setSelectedStorageEntry(primaryStorageEntry);
+
+        assertThat((StorageEntry) mController.mSpinnerPreference.getSelectedItem())
+                .isEqualTo(primaryStorageEntry);
+    }
+}
+
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceControllerTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceControllerTest.java
new file mode 100644
index 0000000..6d9155a
--- /dev/null
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/StorageUsageProgressBarPreferenceControllerTest.java
@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.deviceinfo.storage;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import android.app.usage.StorageStatsManager;
+import android.content.Context;
+import android.os.Looper;
+
+import androidx.preference.Preference;
+import androidx.preference.PreferenceManager;
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.settingslib.widget.UsageProgressBarPreference;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.IOException;
+
+@RunWith(AndroidJUnit4.class)
+public class StorageUsageProgressBarPreferenceControllerTest {
+
+    private static final String FAKE_UUID = "95D9-B3A4";
+    private static final long WAIT_TIMEOUT = 10_000L;
+    private static final long FREE_BYTES = 123L;
+    private static final long TOTAL_BYTES = 456L;
+    private static final long USAGE_BYTES = TOTAL_BYTES - FREE_BYTES;
+
+    private Context mContext;
+    private FakeStorageUsageProgressBarPreferenceController mController;
+    private PreferenceScreen mPreferenceScreen;
+    @Mock
+    private StorageStatsManager mStorageStatsManager;
+
+    @Before
+    public void setUp() throws Exception {
+        if (Looper.myLooper() == null) {
+            Looper.prepare();
+        }
+        MockitoAnnotations.initMocks(this);
+        mContext = spy(ApplicationProvider.getApplicationContext());
+        when(mContext.getSystemService(StorageStatsManager.class)).thenReturn(mStorageStatsManager);
+        mController = new FakeStorageUsageProgressBarPreferenceController(mContext, "key");
+        final PreferenceManager preferenceManager = new PreferenceManager(mContext);
+        mPreferenceScreen = preferenceManager.createPreferenceScreen(mContext);
+        final UsageProgressBarPreference usageProgressBarPreference =
+                new UsageProgressBarPreference(mContext);
+        usageProgressBarPreference.setKey(mController.getPreferenceKey());
+        mPreferenceScreen.addPreference(usageProgressBarPreference);
+    }
+
+    @Test
+    public void setSelectedStorageEntry_primaryStorage_getPrimaryStorageBytes() throws IOException {
+        final StorageEntry defaultInternalStorageEntry =
+                StorageEntry.getDefaultInternalStorageEntry(mContext);
+        when(mStorageStatsManager.getTotalBytes(defaultInternalStorageEntry.getFsUuid()))
+                .thenReturn(TOTAL_BYTES);
+        when(mStorageStatsManager.getFreeBytes(defaultInternalStorageEntry.getFsUuid()))
+                .thenReturn(FREE_BYTES);
+        mController.displayPreference(mPreferenceScreen);
+
+        synchronized (mController.mLock) {
+            mController.setSelectedStorageEntry(defaultInternalStorageEntry);
+            mController.waitUpdateState(WAIT_TIMEOUT);
+        }
+
+        assertThat(mController.mUsedBytes).isEqualTo(USAGE_BYTES);
+        assertThat(mController.mTotalBytes).isEqualTo(TOTAL_BYTES);
+    }
+
+    private class FakeStorageUsageProgressBarPreferenceController
+            extends StorageUsageProgressBarPreferenceController {
+        private final Object mLock = new Object();
+
+        FakeStorageUsageProgressBarPreferenceController(Context context, String key) {
+            super(context, key);
+        }
+
+        @Override
+        public void updateState(Preference preference) {
+            super.updateState(preference);
+            try {
+                mLock.notifyAll();
+            } catch (IllegalMonitorStateException e) {
+                // Catch it for displayPreference to prevent exception by object not locked by
+                // thread before notify. Do nothing.
+            }
+        }
+
+        public void waitUpdateState(long timeout) {
+            try {
+                mLock.wait(timeout);
+            } catch (InterruptedException e) {
+                // Do nothing.
+            }
+        }
+    }
+}
+
diff --git a/tests/unit/src/com/android/settings/deviceinfo/storage/VolumeSizesLoaderTest.java b/tests/unit/src/com/android/settings/deviceinfo/storage/VolumeSizesLoaderTest.java
index 79c5db8..77fd963 100644
--- a/tests/unit/src/com/android/settings/deviceinfo/storage/VolumeSizesLoaderTest.java
+++ b/tests/unit/src/com/android/settings/deviceinfo/storage/VolumeSizesLoaderTest.java
@@ -34,8 +34,10 @@
 @RunWith(AndroidJUnit4.class)
 public class VolumeSizesLoaderTest {
     @Test
-    public void getVolumeSize_getsValidSizes() throws Exception {
+    public void getVolumeSize_privateMountedVolume_getsValidSizes() throws Exception {
         VolumeInfo info = mock(VolumeInfo.class);
+        when(info.getType()).thenReturn(VolumeInfo.TYPE_PRIVATE);
+        when(info.getState()).thenReturn(VolumeInfo.STATE_MOUNTED);
         StorageVolumeProvider storageVolumeProvider = mock(StorageVolumeProvider.class);
         when(storageVolumeProvider.getTotalBytes(any(), any())).thenReturn(10000L);
         when(storageVolumeProvider.getFreeBytes(any(), any())).thenReturn(1000L);
@@ -46,4 +48,19 @@
         assertThat(storageInfo.freeBytes).isEqualTo(1000L);
         assertThat(storageInfo.totalBytes).isEqualTo(10000L);
     }
+
+    @Test
+    public void getVolumeSize_unmountedVolume_getsValidSizes() throws Exception {
+        VolumeInfo info = mock(VolumeInfo.class);
+        when(info.getState()).thenReturn(VolumeInfo.STATE_UNMOUNTED);
+        StorageVolumeProvider storageVolumeProvider = mock(StorageVolumeProvider.class);
+        when(storageVolumeProvider.getTotalBytes(any(), any())).thenReturn(10000L);
+        when(storageVolumeProvider.getFreeBytes(any(), any())).thenReturn(1000L);
+
+        PrivateStorageInfo storageInfo =
+                VolumeSizesLoader.getVolumeSize(storageVolumeProvider, null, info);
+
+        assertThat(storageInfo.freeBytes).isEqualTo(0L);
+        assertThat(storageInfo.totalBytes).isEqualTo(0L);
+    }
 }