Automatically switch to native density for ThemeHostTest

Bug: 128574092
Test: atest android.theme.cts.ThemeHostTest#testThemes

Change-Id: I839943015f81b11bc2b03246ad7869a9ebc43ecb
Merged-In: Icbdcaeb977a8c5b594146f6f6943d855dce43e74
diff --git a/hostsidetests/theme/src/android/theme/cts/ThemeHostTest.java b/hostsidetests/theme/src/android/theme/cts/ThemeHostTest.java
index f5891f4..4dc56f8 100644
--- a/hostsidetests/theme/src/android/theme/cts/ThemeHostTest.java
+++ b/hostsidetests/theme/src/android/theme/cts/ThemeHostTest.java
@@ -83,12 +83,15 @@
 
     private ExecutorCompletionService<Pair<String, File>> mCompletionService;
 
+    // Density to which the device should be restored, or -1 if unnecessary.
+    private int mRestoreDensity;
 
     @Override
     protected void setUp() throws Exception {
         super.setUp();
 
         mDevice = getDevice();
+        mRestoreDensity = resetDensityIfNeeded(mDevice);
         mDevice.executeShellCommand("settings put system font_scale 1.0");
         final String density = getDensityBucketForDevice(mDevice);
         final String referenceZipAssetPath = String.format("/%s.zip", density);
@@ -139,6 +142,8 @@
         // Remove generated images.
         mDevice.executeShellCommand(CLEAR_GENERATED_CMD);
 
+        restoreDensityIfNeeded(mDevice, mRestoreDensity);
+
         super.tearDown();
     }
 
@@ -268,14 +273,26 @@
         return bucket;
     }
 
-    private static int getDensityForDevice(ITestDevice device) throws DeviceNotAvailableException {
+    private static int resetDensityIfNeeded(ITestDevice device) throws DeviceNotAvailableException {
         final String output = device.executeShellCommand(WM_DENSITY);
         final Pattern p = Pattern.compile("Override density: (\\d+)");
         final Matcher m = p.matcher(output);
         if (m.find()) {
-            throw new RuntimeException("Cannot test device running at non-default density: "
-                    + Integer.parseInt(m.group(1)));
+            device.executeShellCommand(WM_DENSITY + " reset");
+            int restoreDensity = Integer.parseInt(m.group(1));
+            return restoreDensity;
         }
+        return -1;
+    }
+
+    private static void restoreDensityIfNeeded(ITestDevice device, int restoreDensity)
+            throws DeviceNotAvailableException {
+        if (restoreDensity > 0) {
+            device.executeShellCommand(WM_DENSITY + " " + restoreDensity);
+        }
+    }
+
+    private static int getDensityForDevice(ITestDevice device) throws DeviceNotAvailableException {
 
         final String densityProp;
         if (device.getSerialNumber().startsWith("emulator-")) {