diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index 45a7999..2752340 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -18,6 +18,7 @@
 
 import static android.view.WindowManager.LayoutParams.FIRST_APPLICATION_WINDOW;
 import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
+import static com.android.server.wm.TaskSnapshotController.*;
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
@@ -72,13 +73,15 @@
     }
 
     @Test
-    public void testSnapshotsDisabled() throws Exception {
+    public void testGetSnapshotMode() throws Exception {
         final WindowState disabledWindow = createWindow(null,
                 FIRST_APPLICATION_WINDOW, sDisplayContent, "disabledWindow");
         disabledWindow.mAppToken.setDisablePreviewSnapshots(true);
-        assertFalse(sWm.mTaskSnapshotController.canSnapshotTask(disabledWindow.getTask()));
+        assertEquals(SNAPSHOT_MODE_APP_THEME,
+                sWm.mTaskSnapshotController.getSnapshotMode(disabledWindow.getTask()));
         final WindowState normalWindow = createWindow(null,
                 FIRST_APPLICATION_WINDOW, sDisplayContent, "normalWindow");
-        assertTrue(sWm.mTaskSnapshotController.canSnapshotTask(normalWindow.getTask()));
+        assertEquals(SNAPSHOT_MODE_REAL,
+                sWm.mTaskSnapshotController.getSnapshotMode(normalWindow.getTask()));
     }
 }