diff --git a/testing/xts/src/com/android/timezone/xts/TimeZoneUpdateHostTest.java b/testing/xts/src/com/android/timezone/xts/TimeZoneUpdateHostTest.java
index 68fabc1..85c6019 100644
--- a/testing/xts/src/com/android/timezone/xts/TimeZoneUpdateHostTest.java
+++ b/testing/xts/src/com/android/timezone/xts/TimeZoneUpdateHostTest.java
@@ -20,13 +20,24 @@
 import com.android.tradefed.config.Option;
 import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.log.LogUtil;
-import com.android.tradefed.testtype.DeviceTestCase;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
 import com.android.tradefed.testtype.IBuildReceiver;
+import com.android.tradefed.testtype.IDeviceTest;
 import com.android.tradefed.util.FileUtil;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
 import java.io.File;
 import java.util.function.BooleanSupplier;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
 /**
  * Class for host-side tests that the time zone rules update feature works as intended. This is
  * intended to give confidence to OEMs that they have implemented / configured the OEM parts of the
@@ -49,9 +60,8 @@
  * This test attempts to handle both of these cases.
  *
  */
-// TODO(nfuller): Switch this to JUnit4 when HostTest supports @Option with JUnit4.
-// http://b/64015928
-public class TimeZoneUpdateHostTest extends DeviceTestCase implements IBuildReceiver {
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class TimeZoneUpdateHostTest implements IDeviceTest, IBuildReceiver {
 
     // These must match equivalent values in RulesManagerService dumpsys code.
     private static final String STAGED_OPERATION_NONE = "None";
@@ -60,6 +70,7 @@
     private static final String INSTALL_STATE_INSTALLED = "Installed";
 
     private IBuildInfo mBuildInfo;
+    private ITestDevice mDevice;
     private File mTempDir;
 
     @Option(name = "oem-data-app-package-name",
@@ -90,28 +101,34 @@
     }
 
     @Override
+    public void setDevice(ITestDevice device) {
+        mDevice = device;
+    }
+
+    @Override
+    public ITestDevice getDevice() {
+        return mDevice;
+    }
+
+    @Before
     public void setUp() throws Exception {
-        super.setUp();
         createTempDir();
         resetDeviceToClean();
     }
 
-    @Override
-    protected void tearDown() throws Exception {
+    @After
+    public void tearDown() throws Exception {
         resetDeviceToClean();
         deleteTempDir();
-        super.tearDown();
     }
 
-    // @Before
-    public void createTempDir() throws Exception {
+    private void createTempDir() throws Exception {
         mTempDir = File.createTempFile("timeZoneUpdateTest", null);
         assertTrue(mTempDir.delete());
         assertTrue(mTempDir.mkdir());
     }
 
-    // @After
-    public void deleteTempDir() throws Exception {
+    private void deleteTempDir() throws Exception {
         FileUtil.recursiveDelete(mTempDir);
     }
 
@@ -119,9 +136,7 @@
      * Reset the device to having no installed time zone data outside of the /system/priv-app
      * version that came with the system image.
      */
-    // @Before
-    // @After
-    public void resetDeviceToClean() throws Exception {
+    private void resetDeviceToClean() throws Exception {
         // If this fails the data app isn't present on device. No point in starting.
         assertTrue(getTimeZoneDataPackageName() + " not installed",
                 isPackageInstalled(getTimeZoneDataPackageName()));
@@ -160,7 +175,7 @@
         assertEquals(STAGED_OPERATION_NONE, getStagedOperationType());
     }
 
-    // @Test
+    @Test
     public void testInstallNewerRulesVersion() throws Exception {
         // This information must match the rules version in test1: IANA version=2030a, revision=1
         String test1VersionInfo = "2030a,1";
@@ -186,7 +201,7 @@
         assertEquals(test1VersionInfo, getCurrentInstalledVersion());
     }
 
-    // @Test
+    @Test
     public void testInstallNewerRulesVersion_secondaryUser() throws Exception {
         ITestDevice device = getDevice();
         if (!device.isMultiUserSupported()) {
@@ -228,7 +243,7 @@
         }
     }
 
-    // @Test
+    @Test
     public void testInstallOlderRulesVersion() throws Exception {
         File appFile = getTimeZoneDataApkFile("test2");
         getDevice().installPackage(appFile, true /* reinstall */);
@@ -331,7 +346,7 @@
     private static void waitForCondition(BooleanSupplier condition) throws Exception {
         int count = 0;
         boolean lastResult;
-        while (!(lastResult = condition.getAsBoolean()) && count++ < 30) {
+        while (!(lastResult = condition.getAsBoolean()) && count++ < 120) {
             Thread.sleep(1000);
         }
         // Some conditions may not be stable so using the lastResult instead of
@@ -389,9 +404,7 @@
     private File getTimeZoneDataApkFile(String testId) throws Exception {
         CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(mBuildInfo);
         String fileName = getTimeZoneDataApkName(testId);
-
-        // TODO(nfuller): Replace with getTestFile(fileName) when it's available in aosp/master.
-        return new File(buildHelper.getTestsDir(), fileName);
+        return buildHelper.getTestFile(fileName);
     }
 
     private boolean isPackageInstalled(String pkg) throws Exception {