diff --git a/tests/robotests/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java b/tests/unit/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java
similarity index 76%
rename from tests/robotests/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java
rename to tests/unit/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java
index ed7f16b24db..e3affefda55 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java
+++ b/tests/unit/src/com/android/settings/deviceinfo/PrivateVolumeOptionMenuControllerTest.java
@@ -18,30 +18,35 @@ package com.android.settings.deviceinfo;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.Activity;
+import android.content.Context;
+import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.os.storage.VolumeInfo;
 import android.view.Menu;
 import android.view.MenuInflater;
 import android.view.MenuItem;
 
-import com.android.settings.R;
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.settings.testutils.ResourcesUtils;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Answers;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.shadows.ShadowApplication;
 
-@RunWith(RobolectricTestRunner.class)
+@RunWith(AndroidJUnit4.class)
 public class PrivateVolumeOptionMenuControllerTest {
 
     @Mock
@@ -57,11 +62,13 @@ public class PrivateVolumeOptionMenuControllerTest {
     @Mock
     private VolumeInfo mPrimaryInfo;
 
+    private Context mContext;
     private PrivateVolumeOptionMenuController mController;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+        mContext = spy(ApplicationProvider.getApplicationContext());
 
         when(mVolumeInfo.getType()).thenReturn(VolumeInfo.TYPE_PRIVATE);
         when(mVolumeInfo.isMountedWritable()).thenReturn(true);
@@ -69,15 +76,15 @@ public class PrivateVolumeOptionMenuControllerTest {
         when(mMenu.findItem(anyInt())).thenReturn(mMigrateMenuItem);
         when(mMigrateMenuItem.getItemId()).thenReturn(100);
 
-        mController = new PrivateVolumeOptionMenuController(
-                Robolectric.setupActivity(Activity.class), mPrimaryInfo, mPm);
+        mController = new PrivateVolumeOptionMenuController(mContext, mPrimaryInfo, mPm);
     }
 
     @Test
     public void testMigrateDataMenuItemIsAdded() {
         mController.onCreateOptionsMenu(mMenu, mMenuInflater);
 
-        verify(mMenu).add(Menu.NONE, 100, Menu.NONE, R.string.storage_menu_migrate);
+        verify(mMenu).add(Menu.NONE, 100, Menu.NONE, ResourcesUtils.getResourcesId(
+                mContext, "string", "storage_menu_migrate"));
     }
 
     @Test
@@ -115,14 +122,18 @@ public class PrivateVolumeOptionMenuControllerTest {
     @Test
     public void testMigrateDataGoesToMigrateWizard() {
         when(mPm.getPrimaryStorageCurrentVolume()).thenReturn(mVolumeInfo);
+        doNothing().when(mContext).startActivity(any(Intent.class));
 
         mController.onCreateOptionsMenu(mMenu, mMenuInflater);
         mController.onPrepareOptionsMenu(mMenu);
 
         assertThat(mController.onOptionsItemSelected(mMigrateMenuItem)).isTrue();
-        ShadowApplication shadowApplication = ShadowApplication.getInstance();
-        assertThat(shadowApplication).isNotNull();
-        assertThat(shadowApplication.getNextStartedActivity().getComponent().getClassName())
+
+        final ArgumentCaptor<Intent> intentCaptor = ArgumentCaptor.forClass(
+                Intent.class);
+        verify(mContext).startActivity(intentCaptor.capture());
+        final Intent startedIntent = intentCaptor.getValue();
+        assertThat(startedIntent.getComponent().getClassName())
                 .isEqualTo(StorageWizardMigrateConfirm.class.getName());
     }
 }