Merge "Snap for 7562097 from 879edbea786f7670c33640e480281f3cfab5d1d2 to sdk-release" into sdk-release
diff --git a/common_util/com/android/tradefed/result/error/TestErrorIdentifier.java b/common_util/com/android/tradefed/result/error/TestErrorIdentifier.java
index 59bdd07..4931c25 100644
--- a/common_util/com/android/tradefed/result/error/TestErrorIdentifier.java
+++ b/common_util/com/android/tradefed/result/error/TestErrorIdentifier.java
@@ -28,7 +28,8 @@
     TEST_ABORTED(530_005, FailureStatus.TEST_FAILURE),
     OUTPUT_PARSER_ERROR(530_006, FailureStatus.TEST_FAILURE),
     TEST_BINARY_EXIT_CODE_ERROR(530_007, FailureStatus.TEST_FAILURE),
-    TEST_BINARY_TIMED_OUT(530_008, FailureStatus.TIMED_OUT);
+    TEST_BINARY_TIMED_OUT(530_008, FailureStatus.TIMED_OUT),
+    MODIFIED_FOLDABLE_STATE(530_009, FailureStatus.TEST_FAILURE);
 
     private final long code;
     private final @Nonnull FailureStatus status;
diff --git a/device_build_interfaces/com/android/tradefed/device/DeviceFoldableState.java b/device_build_interfaces/com/android/tradefed/device/DeviceFoldableState.java
new file mode 100644
index 0000000..873c2c5
--- /dev/null
+++ b/device_build_interfaces/com/android/tradefed/device/DeviceFoldableState.java
@@ -0,0 +1,77 @@
+/*
+ * 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.tradefed.device;
+
+/**
+ * Representation of device foldable state as returned by "cmd device_state print-states".
+ */
+public class DeviceFoldableState implements Comparable<DeviceFoldableState> {
+
+    private final long mIdentifier;
+    private final String mName;
+
+    public DeviceFoldableState(long identifier, String name) {
+        mIdentifier = identifier;
+        mName = name;
+    }
+
+    public long getIdentifier() {
+        return mIdentifier;
+    }
+
+    @Override
+    public String toString() {
+        return "foldable:" + mIdentifier + ":" + mName;
+    }
+
+    @Override
+    public int hashCode() {
+        final int prime = 31;
+        int result = 1;
+        result = prime * result + (int) (mIdentifier ^ (mIdentifier >>> 32));
+        result = prime * result + ((mName == null) ? 0 : mName.hashCode());
+        return result;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj)
+            return true;
+        if (obj == null)
+            return false;
+        if (getClass() != obj.getClass())
+            return false;
+        DeviceFoldableState other = (DeviceFoldableState) obj;
+        if (mIdentifier != other.mIdentifier)
+            return false;
+        if (mName == null) {
+            if (other.mName != null)
+                return false;
+        } else if (!mName.equals(other.mName))
+            return false;
+        return true;
+    }
+
+    @Override
+    public int compareTo(DeviceFoldableState o) {
+        if (this.mIdentifier == o.mIdentifier) {
+            return 0;
+        } else if (this.mIdentifier < o.mIdentifier) {
+            return 1;
+        }
+        return -1;
+    }
+}
diff --git a/device_build_interfaces/com/android/tradefed/device/ITestDevice.java b/device_build_interfaces/com/android/tradefed/device/ITestDevice.java
index 67860cb..bdd0986 100644
--- a/device_build_interfaces/com/android/tradefed/device/ITestDevice.java
+++ b/device_build_interfaces/com/android/tradefed/device/ITestDevice.java
@@ -977,4 +977,19 @@
      * @throws DeviceNotAvailableException
      */
     public Set<Long> listDisplayIds() throws DeviceNotAvailableException;
+
+    /**
+     * Returns the list of foldable states on the device. Can be obtained with
+     * "cmd device_state print-states".
+     *
+     * @throws DeviceNotAvailableException
+     */
+    public Set<DeviceFoldableState> getFoldableStates() throws DeviceNotAvailableException;
+
+    /**
+     * Returns the current foldable state of the device or null if some issues occurred.
+     *
+     * @throws DeviceNotAvailableException
+     */
+    public DeviceFoldableState getCurrentFoldableState() throws DeviceNotAvailableException;
 }
diff --git a/javatests/com/android/tradefed/command/CommandSchedulerTest.java b/javatests/com/android/tradefed/command/CommandSchedulerTest.java
index 3a11b0f..12bfccb 100644
--- a/javatests/com/android/tradefed/command/CommandSchedulerTest.java
+++ b/javatests/com/android/tradefed/command/CommandSchedulerTest.java
@@ -15,13 +15,22 @@
  */
 package com.android.tradefed.command;
 
-import static org.easymock.EasyMock.getCurrentArguments;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.android.ddmlib.IDevice;
 import com.android.tradefed.command.CommandFileParser.CommandLine;
@@ -62,14 +71,15 @@
 import com.android.tradefed.util.keystore.DryRunKeyStore;
 import com.android.tradefed.util.keystore.IKeyStoreClient;
 
-import org.easymock.EasyMock;
-import org.easymock.IAnswer;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.AdditionalMatchers;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -82,23 +92,23 @@
 import java.util.List;
 import java.util.Map;
 
-
 /** Unit tests for {@link CommandScheduler}. */
 @RunWith(JUnit4.class)
 public class CommandSchedulerTest {
 
     private CommandScheduler mScheduler;
-    private ITestInvocation mMockInvocation;
-    private MockDeviceManager mMockManager;
-    private IConfigurationFactory mMockConfigFactory;
-    private IConfiguration mMockConfiguration;
+    @Mock ITestInvocation mMockInvocation;
+    private MockDeviceManager mFakeDeviceManager;
+    @Mock IConfigurationFactory mMockConfigFactory;
+    @Mock IConfiguration mMockConfiguration;
     private CommandOptions mCommandOptions;
     private DeviceSelectionOptions mDeviceOptions;
-    private CommandFileParser mMockCmdFileParser;
-    private List<IDeviceConfiguration> mMockDeviceConfig;
-    private ConfigurationDescriptor mMockConfigDescriptor;
-    private IKeyStoreClient mMockKeyStoreClient;
+    private CommandFileParser mCommandFileParser;
+    private List<IDeviceConfiguration> mDeviceConfigList;
+    private ConfigurationDescriptor mConfigDescriptor;
+    @Mock IKeyStoreClient mMockKeyStoreClient;
     private IInvocationContext mContext;
+    private boolean mIsFirstInvoke = true; // For testRun_rescheduled()
 
     class TestableCommandScheduler extends CommandScheduler {
 
@@ -109,7 +119,7 @@
 
         @Override
         protected IDeviceManager getDeviceManager() {
-            return mMockManager;
+            return mFakeDeviceManager;
         }
 
         @Override
@@ -149,7 +159,7 @@
 
         @Override
         CommandFileParser createCommandFileParser() {
-            return mMockCmdFileParser;
+            return mCommandFileParser;
         }
 
         @Override
@@ -160,17 +170,14 @@
 
     @Before
     public void setUp() throws Exception {
-        mMockInvocation = EasyMock.createMock(ITestInvocation.class);
-        EasyMock.expect(mMockInvocation.getExitInfo()).andStubReturn(new ExitInformation());
-        mMockManager = new MockDeviceManager(0);
-        mMockConfigFactory = EasyMock.createMock(IConfigurationFactory.class);
-        mMockKeyStoreClient = EasyMock.createMock(IKeyStoreClient.class);
-        mMockConfiguration = EasyMock.createMock(IConfiguration.class);
-        EasyMock.expect(mMockConfiguration.getTests()).andStubReturn(new ArrayList<>());
-        EasyMock.expect(
-                        mMockConfiguration.getConfigurationObject(
-                                ProxyConfiguration.PROXY_CONFIG_TYPE_KEY))
-                .andStubReturn(null);
+        MockitoAnnotations.initMocks(this);
+
+        when(mMockInvocation.getExitInfo()).thenReturn(new ExitInformation());
+
+        mFakeDeviceManager = new MockDeviceManager(0);
+        when(mMockConfiguration.getTests()).thenReturn(new ArrayList<>());
+        when(mMockConfiguration.getConfigurationObject(ProxyConfiguration.PROXY_CONFIG_TYPE_KEY))
+                .thenReturn(null);
         mCommandOptions = new CommandOptions();
         // Avoid any issue related to env. variable.
         mDeviceOptions =
@@ -180,8 +187,8 @@
                         return null;
                     }
                 };
-        mMockDeviceConfig = new ArrayList<IDeviceConfiguration>();
-        mMockConfigDescriptor = new ConfigurationDescriptor();
+        mDeviceConfigList = new ArrayList<IDeviceConfiguration>();
+        mConfigDescriptor = new ConfigurationDescriptor();
         mContext = new InvocationContext();
 
         mScheduler = new TestableCommandScheduler();
@@ -195,34 +202,11 @@
         }
     }
 
-    /**
-     * Switch all mock objects to replay mode
-     */
-    private void replayMocks(Object... additionalMocks) {
-        EasyMock.replay(
-                mMockConfigFactory, mMockConfiguration, mMockInvocation, mMockKeyStoreClient);
-        for (Object mock : additionalMocks) {
-            EasyMock.replay(mock);
-        }
-    }
-
-    /**
-     * Verify all mock objects
-     */
-    private void verifyMocks(Object... additionalMocks) {
-        EasyMock.verify(
-                mMockConfigFactory, mMockConfiguration, mMockInvocation, mMockKeyStoreClient);
-        for (Object mock : additionalMocks) {
-            EasyMock.verify(mock);
-        }
-        mMockManager.assertDevicesFreed();
-    }
-
     /** Test {@link CommandScheduler#run()} when no configs have been added */
     @Test
     public void testRun_empty() throws InterruptedException {
-        mMockManager.setNumDevices(1);
-        replayMocks();
+        mFakeDeviceManager.setNumDevices(1);
+
         mScheduler.start();
         while (!mScheduler.isAlive()) {
             Thread.sleep(10);
@@ -230,7 +214,8 @@
         mScheduler.shutdown();
         // expect run not to block
         mScheduler.join();
-        verifyMocks();
+
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /** Test {@link CommandScheduler#addCommand(String[])} when help mode is specified */
@@ -238,30 +223,44 @@
     public void testAddConfig_configHelp() throws ConfigurationException {
         String[] args = new String[] {"test"};
         mCommandOptions.setHelpMode(true);
-        setCreateConfigExpectations(args, 1);
-        // expect
-        mMockConfigFactory.printHelpForConfig(EasyMock.aryEq(args), EasyMock.eq(true),
-                EasyMock.eq(System.out));
-        replayMocks();
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
-        verifyMocks();
+
+        verify(mMockConfigFactory)
+                .printHelpForConfig(AdditionalMatchers.aryEq(args), eq(true), eq(System.out));
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /** Test {@link CommandScheduler#run()} when one config has been added */
     @Test
     public void testRun_oneConfig() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(2);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
-        mMockConfiguration.validateOptions();
-        replayMocks();
+        mFakeDeviceManager.setNumDevices(2);
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -271,16 +270,21 @@
     @Test
     public void testRemoveAllCommands() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(0);
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        replayMocks();
+        mFakeDeviceManager.setNumDevices(0);
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
         assertEquals(1, mScheduler.getAllCommandsSize());
         mScheduler.removeAllCommands();
         assertEquals(0, mScheduler.getAllCommandsSize());
-        verifyMocks();
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /** Test {@link CommandScheduler#run()} when one config has been added in dry-run mode */
@@ -288,17 +292,13 @@
     public void testRun_dryRun() throws Throwable {
         String[] dryRunArgs = new String[] {"--dry-run"};
         mCommandOptions.setDryRunMode(true);
-        mMockManager.setNumDevices(2);
-        setCreateConfigExpectations(dryRunArgs, 1);
+        mFakeDeviceManager.setNumDevices(2);
+        setCreateConfigExpectations(dryRunArgs);
 
         // add a second command, to verify the first dry-run command did not get added
         String[] args2 = new String[] {"test"};
-        setCreateConfigExpectations(args2, 1);
-        setExpectedInvokeCalls(1);
-        mMockConfiguration.validateOptions();
-        EasyMock.expectLastCall().times(2);
+        setCreateConfigExpectations(args2);
 
-        replayMocks();
         mScheduler.start();
         assertFalse(mScheduler.addCommand(dryRunArgs));
         // the same config object is being used, so clear its state
@@ -306,7 +306,18 @@
         assertTrue(mScheduler.addCommand(args2));
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+        verify(mMockConfiguration, times(2)).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(dryRunArgs), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -325,15 +336,14 @@
                 };
         String[] dryRunArgs =
                 new String[] {"empty", "--noisy-dry-run", "--min-loop-time", "USE_KEYSTORE@fake"};
-        mMockManager.setNumDevices(2);
-        //setCreateConfigExpectations(dryRunArgs, 1);
+        mFakeDeviceManager.setNumDevices(2);
 
-        replayMocks();
         mScheduler.start();
         assertFalse(mScheduler.addCommand(dryRunArgs));
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
+
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -343,30 +353,20 @@
     @Test
     @SuppressWarnings("unchecked")
     public void testExecCommand() throws Throwable {
-        String[] args = new String[] {
-            "foo"
-        };
-        setCreateConfigExpectations(args, 1);
-        mMockInvocation.invoke(
-                (IInvocationContext) EasyMock.anyObject(),
-                (IConfiguration) EasyMock.anyObject(),
-                (IRescheduler) EasyMock.anyObject(),
-                (ITestInvocationListener) EasyMock.anyObject(),
-                EasyMock.anyObject());
-        EasyMock.expectLastCall().times(1);
-        mMockConfiguration.validateOptions();
-        IDevice mockIDevice = EasyMock.createMock(IDevice.class);
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        EasyMock.expect(mockDevice.getSerialNumber()).andStubReturn("serial");
-        EasyMock.expect(mockDevice.getDeviceState()).andStubReturn(TestDeviceState.ONLINE);
-        mockDevice.setRecoveryMode(EasyMock.eq(RecoveryMode.AVAILABLE));
-        EasyMock.expect(mockDevice.getIDevice()).andStubReturn(mockIDevice);
-        IScheduledInvocationListener mockListener = EasyMock
-                .createMock(IScheduledInvocationListener.class);
-        mockListener.invocationInitiated((IInvocationContext) EasyMock.anyObject());
-        mockListener.invocationComplete((IInvocationContext)EasyMock.anyObject(),
-                (Map<ITestDevice, FreeDeviceState>)EasyMock.anyObject());
-        EasyMock.expect(mockDevice.waitForDeviceShell(EasyMock.anyLong())).andReturn(true);
+        String[] args = new String[] {"foo"};
+        setCreateConfigExpectations(args);
+
+        IDevice mockIDevice = mock(IDevice.class);
+        ITestDevice mockDevice = mock(ITestDevice.class);
+        when(mockDevice.getSerialNumber()).thenReturn("serial");
+        when(mockDevice.getDeviceState()).thenReturn(TestDeviceState.ONLINE);
+        mockDevice.setRecoveryMode(eq(RecoveryMode.AVAILABLE));
+        when(mockDevice.getIDevice()).thenReturn(mockIDevice);
+        IScheduledInvocationListener mockListener = mock(IScheduledInvocationListener.class);
+        mockListener.invocationInitiated((IInvocationContext) any());
+        mockListener.invocationComplete(
+                (IInvocationContext) any(), (Map<ITestDevice, FreeDeviceState>) any());
+        when(mockDevice.waitForDeviceShell(anyLong())).thenReturn(true);
         mScheduler =
                 new TestableCommandScheduler() {
                     @Override
@@ -380,26 +380,24 @@
                         return results;
                     }
                 };
-        replayMocks(mockDevice, mockListener);
+
         mScheduler.start();
         mScheduler.execCommand(mockListener, args);
         mScheduler.shutdownOnEmpty();
-        mScheduler.join(2*1000);
-        verifyMocks(mockListener);
-    }
+        mScheduler.join(2 * 1000);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any(),
+                        any());
+        verify(mMockConfiguration).validateOptions();
 
-    /**
-     * Sets the number of expected
-     * {@link ITestInvocation#invoke(IInvocationContext, IConfiguration, IRescheduler,
-     *      ITestInvocationListener[])} calls
-     *
-     * @param times
-     */
-    private void setExpectedInvokeCalls(int times) throws Throwable {
-        mMockInvocation.invoke((IInvocationContext)EasyMock.anyObject(),
-                (IConfiguration)EasyMock.anyObject(), (IRescheduler)EasyMock.anyObject(),
-                (ITestInvocationListener)EasyMock.anyObject());
-        EasyMock.expectLastCall().times(times);
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -410,23 +408,31 @@
     @Test
     public void testRun_configSerial() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(2);
-        setCreateConfigExpectations(args, 2);
+        mFakeDeviceManager.setNumDevices(2);
+        setCreateConfigExpectations(args);
         // allocate and free a device to get its serial
-        ITestDevice dev = mMockManager.allocateDevice();
+        ITestDevice dev = mFakeDeviceManager.allocateDevice();
         mDeviceOptions.addSerial(dev.getSerialNumber());
-        setExpectedInvokeCalls(1);
         mMockConfiguration.validateOptions();
-        mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.addCommand(args);
-        mMockManager.freeDevice(dev, FreeDeviceState.AVAILABLE);
+        mFakeDeviceManager.freeDevice(dev, FreeDeviceState.AVAILABLE);
 
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
+        verify(mMockInvocation, times(2))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+        verify(mMockConfiguration, times(3)).validateOptions();
+        verify(mMockConfigFactory, times(2))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -438,67 +444,81 @@
     @Test
     public void testRun_configExcludeSerial() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(2);
-        setCreateConfigExpectations(args, 2);
+        mFakeDeviceManager.setNumDevices(2);
+        setCreateConfigExpectations(args);
         // allocate and free a device to get its serial
-        ITestDevice dev = mMockManager.allocateDevice();
+        ITestDevice dev = mFakeDeviceManager.allocateDevice();
         mDeviceOptions.addExcludeSerial(dev.getSerialNumber());
-        ITestDevice expectedDevice = mMockManager.allocateDevice();
-        setExpectedInvokeCalls(1);
+        ITestDevice expectedDevice = mFakeDeviceManager.allocateDevice();
         mMockConfiguration.validateOptions();
-        mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.addCommand(args);
-        mMockManager.freeDevice(dev, FreeDeviceState.AVAILABLE);
-        mMockManager.freeDevice(expectedDevice, FreeDeviceState.AVAILABLE);
+        mFakeDeviceManager.freeDevice(dev, FreeDeviceState.AVAILABLE);
+        mFakeDeviceManager.freeDevice(expectedDevice, FreeDeviceState.AVAILABLE);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
+        verify(mMockInvocation, times(2))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+        verify(mMockConfiguration, times(3)).validateOptions();
+        verify(mMockConfigFactory, times(2))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /** Test {@link CommandScheduler#run()} when one config has been rescheduled */
     @Test
     public void testRun_rescheduled() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(2);
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        final IConfiguration rescheduledConfig = EasyMock.createMock(IConfiguration.class);
-        EasyMock.expect(rescheduledConfig.getCommandOptions()).andStubReturn(mCommandOptions);
-        EasyMock.expect(rescheduledConfig.getDeviceRequirements()).andStubReturn(
-                mDeviceOptions);
-        EasyMock.expect(rescheduledConfig.getDeviceConfig()).andStubReturn(mMockDeviceConfig);
-        EasyMock.expect(rescheduledConfig.getCommandLine()).andStubReturn("");
-        EasyMock.expect(rescheduledConfig.getConfigurationDescription())
-            .andStubReturn(mMockConfigDescriptor);
+        mFakeDeviceManager.setNumDevices(2);
+        setCreateConfigExpectations(args);
 
-        // an ITestInvocationn#invoke response for calling reschedule
-        IAnswer<Object> rescheduleAndThrowAnswer = new IAnswer<Object>() {
-            @Override
-            public Object answer() throws Throwable {
-                IRescheduler rescheduler =  (IRescheduler) EasyMock.getCurrentArguments()[2];
-                rescheduler.scheduleConfig(rescheduledConfig);
-                throw new DeviceNotAvailableException("not avail", "fakeserial");
-            }
-        };
+        final IConfiguration mockRescheduledConfig = mock(IConfiguration.class);
+        when(mockRescheduledConfig.getCommandOptions()).thenReturn(mCommandOptions);
+        when(mockRescheduledConfig.getDeviceRequirements()).thenReturn(mDeviceOptions);
+        when(mockRescheduledConfig.getDeviceConfig()).thenReturn(mDeviceConfigList);
+        when(mockRescheduledConfig.getCommandLine()).thenReturn("");
+        when(mockRescheduledConfig.getConfigurationDescription()).thenReturn(mConfigDescriptor);
 
-        mMockInvocation.invoke(EasyMock.<IInvocationContext>anyObject(),
-                EasyMock.<IConfiguration>anyObject(), EasyMock.<IRescheduler>anyObject(),
-                EasyMock.<ITestInvocationListener>anyObject());
-        EasyMock.expectLastCall().andAnswer(rescheduleAndThrowAnswer);
+        // The first call sets recheduler and throws. The second call is successful.
+        doAnswer(
+                        invocation -> {
+                            if (mIsFirstInvoke) {
+                                mIsFirstInvoke = false;
 
-        // expect one more success call
-        setExpectedInvokeCalls(1);
+                                IRescheduler rescheduler =
+                                        (IRescheduler) invocation.getArguments()[2];
+                                rescheduler.scheduleConfig(mockRescheduledConfig);
+                                throw new DeviceNotAvailableException("not avail", "fakeserial");
+                            } else {
+                                return null;
+                            }
+                        })
+                .when(mMockInvocation)
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
 
-        replayMocks(rescheduledConfig);
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
 
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
+        verify(mMockConfiguration).validateOptions();
+        verify(mMockInvocation, times(2))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
     }
 
     /**
@@ -509,19 +529,19 @@
     @Test
     public void testAddCommandFile() throws ConfigurationException {
         // set number of devices to 0 so we can verify command presence
-        mMockManager.setNumDevices(0);
+        mFakeDeviceManager.setNumDevices(0);
         List<String> extraArgs = Arrays.asList("--bar");
-        setCreateConfigExpectations(new String[] {"foo", "--bar"}, 1);
+        setCreateConfigExpectations(new String[] {"foo", "--bar"});
         mMockConfiguration.validateOptions();
-        final List<CommandLine> cmdFileContent = Arrays.asList(new CommandLine(
-                Arrays.asList("foo"), null, 0));
-        mMockCmdFileParser = new CommandFileParser() {
-            @Override
-            public List<CommandLine> parseFile(File cmdFile) {
-                return cmdFileContent;
-            }
-        };
-        replayMocks();
+        final List<CommandLine> cmdFileContent =
+                Arrays.asList(new CommandLine(Arrays.asList("foo"), null, 0));
+        mCommandFileParser =
+                new CommandFileParser() {
+                    @Override
+                    public List<CommandLine> parseFile(File cmdFile) {
+                        return cmdFileContent;
+                    }
+                };
 
         mScheduler.start();
         mScheduler.addCommandFile("mycmd.txt", extraArgs);
@@ -539,35 +559,34 @@
     @Test
     public void testAddCommandFile_reload() throws ConfigurationException {
         // set number of devices to 0 so we can verify command presence
-        mMockManager.setNumDevices(0);
-        String[] addCommandArgs = new String[]{"fromcommand"};
+        mFakeDeviceManager.setNumDevices(0);
+        String[] addCommandArgs = new String[] {"fromcommand"};
         List<String> extraArgs = Arrays.asList("--bar");
 
-        setCreateConfigExpectations(addCommandArgs, 1);
+        setCreateConfigExpectations(addCommandArgs);
         String[] cmdFile1Args = new String[] {"fromFile1", "--bar"};
-        setCreateConfigExpectations(cmdFile1Args, 1);
+        setCreateConfigExpectations(cmdFile1Args);
         String[] cmdFile2Args = new String[] {"fromFile2", "--bar"};
-        setCreateConfigExpectations(cmdFile2Args, 1);
+        setCreateConfigExpectations(cmdFile2Args);
 
-        mMockConfiguration.validateOptions();
-        EasyMock.expectLastCall().times(3);
+        final List<CommandLine> cmdFileContent1 =
+                Arrays.asList(new CommandLine(Arrays.asList("fromFile1"), null, 0));
+        final List<CommandLine> cmdFileContent2 =
+                Arrays.asList(new CommandLine(Arrays.asList("fromFile2"), null, 0));
+        mCommandFileParser =
+                new CommandFileParser() {
+                    boolean firstCall = true;
 
-        final List<CommandLine> cmdFileContent1 = Arrays.asList(new CommandLine(
-                Arrays.asList("fromFile1"), null, 0));
-        final List<CommandLine> cmdFileContent2 = Arrays.asList(new CommandLine(
-                Arrays.asList("fromFile2"), null, 0));
-        mMockCmdFileParser = new CommandFileParser() {
-            boolean firstCall = true;
-            @Override
-            public List<CommandLine> parseFile(File cmdFile) {
-                if (firstCall) {
-                    firstCall = false;
-                    return cmdFileContent1;
-                }
-                return cmdFileContent2;
-            }
-        };
-        replayMocks();
+                    @Override
+                    public List<CommandLine> parseFile(File cmdFile) {
+                        if (firstCall) {
+                            firstCall = false;
+                            return cmdFileContent1;
+                        }
+                        return cmdFileContent2;
+                    }
+                };
+
         mScheduler.start();
         mScheduler.setCommandFileReload(true);
         mScheduler.addCommand(addCommandArgs);
@@ -587,28 +606,28 @@
         Collections.sort(cmds, new CommandTrackerIdComparator());
         Assert.assertArrayEquals(addCommandArgs, cmds.get(0).getArgs());
         Assert.assertArrayEquals(cmdFile2Args, cmds.get(1).getArgs());
+        verify(mMockConfiguration, times(3)).validateOptions();
     }
 
     /** Verify attempts to add the same commmand file in reload mode are rejected */
     @Test
     public void testAddCommandFile_twice() throws ConfigurationException {
         // set number of devices to 0 so we can verify command presence
-        mMockManager.setNumDevices(0);
+        mFakeDeviceManager.setNumDevices(0);
         String[] cmdFile1Args = new String[] {"fromFile1"};
-        setCreateConfigExpectations(cmdFile1Args, 1);
-        setCreateConfigExpectations(cmdFile1Args, 1);
-        mMockConfiguration.validateOptions();
-        EasyMock.expectLastCall().times(2);
+        setCreateConfigExpectations(cmdFile1Args);
+        setCreateConfigExpectations(cmdFile1Args);
 
-        final List<CommandLine> cmdFileContent1 = Arrays.asList(new CommandLine(
-                Arrays.asList("fromFile1"), null, 0));
-        mMockCmdFileParser = new CommandFileParser() {
-            @Override
-            public List<CommandLine> parseFile(File cmdFile) {
-                return cmdFileContent1;
-            }
-        };
-        replayMocks();
+        final List<CommandLine> cmdFileContent1 =
+                Arrays.asList(new CommandLine(Arrays.asList("fromFile1"), null, 0));
+        mCommandFileParser =
+                new CommandFileParser() {
+                    @Override
+                    public List<CommandLine> parseFile(File cmdFile) {
+                        return cmdFileContent1;
+                    }
+                };
+
         mScheduler.start();
         mScheduler.setCommandFileReload(true);
         mScheduler.addCommandFile("mycmd.txt", Collections.<String>emptyList());
@@ -625,12 +644,13 @@
         cmds = mScheduler.getCommandTrackers();
         assertEquals(1, cmds.size());
         Assert.assertArrayEquals(cmdFile1Args, cmds.get(0).getArgs());
+        verify(mMockConfiguration, times(2)).validateOptions();
     }
 
     /** Test {@link CommandScheduler#shutdown()} when no devices are available. */
     @Test
     public void testShutdown() throws Exception {
-        mMockManager.setNumDevices(0);
+        mFakeDeviceManager.setNumDevices(0);
         mScheduler.start();
         while (!mScheduler.isAlive()) {
             Thread.sleep(10);
@@ -642,30 +662,23 @@
         // test will hang if not successful
     }
 
-    /**
-     * Set EasyMock expectations for a create configuration call.
-     */
-    private void setCreateConfigExpectations(String[] args, int times)
-            throws ConfigurationException {
+    /** Set EasyMock expectations for a create configuration call. */
+    private void setCreateConfigExpectations(String[] args) throws ConfigurationException {
         List<String> nullArg = null;
-        EasyMock.expect(
-                mMockConfigFactory.createConfigurationFromArgs(EasyMock.aryEq(args),
-                EasyMock.eq(nullArg), (IKeyStoreClient)EasyMock.anyObject()))
-                .andReturn(mMockConfiguration)
-                .times(times);
-        EasyMock.expect(mMockConfiguration.getCommandOptions()).andStubReturn(mCommandOptions);
-        EasyMock.expect(mMockConfiguration.getDeviceRequirements()).andStubReturn(
-                mDeviceOptions);
-        EasyMock.expect(mMockConfiguration.getDeviceConfig()).andStubReturn(mMockDeviceConfig);
-        EasyMock.expect(mMockConfiguration.getCommandLine()).andStubReturn("");
-        EasyMock.expect(mMockConfiguration.getConfigurationDescription())
-            .andStubReturn(mMockConfigDescriptor);
+        when(mMockConfigFactory.createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any()))
+                .thenReturn(mMockConfiguration);
+        when(mMockConfiguration.getCommandOptions()).thenReturn(mCommandOptions);
+        when(mMockConfiguration.getDeviceRequirements()).thenReturn(mDeviceOptions);
+        when(mMockConfiguration.getDeviceConfig()).thenReturn(mDeviceConfigList);
+        when(mMockConfiguration.getCommandLine()).thenReturn("");
+        when(mMockConfiguration.getConfigurationDescription()).thenReturn(mConfigDescriptor);
 
         // Assume all legacy test are single device
-        if (mMockDeviceConfig.isEmpty()) {
+        if (mDeviceConfigList.isEmpty()) {
             IDeviceConfiguration mockConfig = new DeviceConfigurationHolder("device");
             mockConfig.addSpecificConfig(mDeviceOptions);
-            mMockDeviceConfig.add(mockConfig);
+            mDeviceConfigList.add(mockConfig);
         }
     }
 
@@ -673,18 +686,28 @@
     @Test
     public void testDeviceReleased() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(1);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
-        mMockConfiguration.validateOptions();
-        replayMocks();
+        mFakeDeviceManager.setNumDevices(1);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        verifyMocks();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
     }
 
     /**
@@ -694,40 +717,33 @@
     @Test
     public void testDeviceReleasedEarly() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(1);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        setCreateConfigExpectations(args, 2);
+        mFakeDeviceManager.setNumDevices(1);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        setCreateConfigExpectations(args);
 
-        mMockInvocation.invoke(
-                (IInvocationContext) EasyMock.anyObject(),
-                (IConfiguration) EasyMock.anyObject(),
-                (IRescheduler) EasyMock.anyObject(),
-                (ITestInvocationListener) EasyMock.anyObject());
-        EasyMock.expectLastCall()
-                .andAnswer(
-                        new IAnswer<Object>() {
-                            @Override
-                            public Object answer() throws Throwable {
-                                IInvocationContext context =
-                                        (IInvocationContext) getCurrentArguments()[0];
-                                IScheduledInvocationListener listener =
-                                        (IScheduledInvocationListener) getCurrentArguments()[3];
-                                Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
-                                for (ITestDevice device : context.getDevices()) {
-                                    deviceStates.put(device, FreeDeviceState.AVAILABLE);
-                                }
-                                context.markReleasedEarly();
-                                listener.releaseDevices(context, deviceStates);
-                                RunUtil.getDefault().sleep(500);
-                                return null;
+        doAnswer(
+                        invocation -> {
+                            IInvocationContext context =
+                                    (IInvocationContext) invocation.getArguments()[0];
+                            IScheduledInvocationListener listener =
+                                    (IScheduledInvocationListener) invocation.getArguments()[3];
+                            Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
+                            for (ITestDevice device : context.getDevices()) {
+                                deviceStates.put(device, FreeDeviceState.AVAILABLE);
                             }
-                        });
-        // Second invocation runs properly
-        setExpectedInvokeCalls(1);
+                            context.markReleasedEarly();
+                            listener.releaseDevices(context, deviceStates);
+                            RunUtil.getDefault().sleep(500);
+                            return null;
+                        })
+                .when(mMockInvocation)
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
 
-        mMockConfiguration.validateOptions();
-        EasyMock.expectLastCall().times(2);
-        replayMocks();
+        // Second invocation runs properly
         mScheduler.start();
         mScheduler.addCommand(args);
         RunUtil.getDefault().sleep(100);
@@ -735,8 +751,18 @@
         RunUtil.getDefault().sleep(200);
         mScheduler.shutdown();
         mScheduler.join();
-        verifyMocks();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockInvocation, times(2))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+        verify(mMockConfiguration, times(2)).validateOptions();
+        verify(mMockConfigFactory, times(2))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
         assertNull(mScheduler.getLastInvocationThrowable());
     }
 
@@ -747,39 +773,34 @@
     @Test
     public void testDeviceReleasedEarly_conflict() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevices(1);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        setCreateConfigExpectations(args, 2);
-        EasyMock.reset(mMockInvocation);
-        EasyMock.expect(mMockInvocation.getExitInfo()).andReturn(null);
-        mMockInvocation.invoke(
-                (IInvocationContext) EasyMock.anyObject(),
-                (IConfiguration) EasyMock.anyObject(),
-                (IRescheduler) EasyMock.anyObject(),
-                (ITestInvocationListener) EasyMock.anyObject());
-        EasyMock.expectLastCall()
-                .andAnswer(
-                        new IAnswer<Object>() {
-                            @Override
-                            public Object answer() throws Throwable {
-                                IInvocationContext context =
-                                        (IInvocationContext) getCurrentArguments()[0];
-                                IScheduledInvocationListener listener =
-                                        (IScheduledInvocationListener) getCurrentArguments()[3];
-                                Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
-                                for (ITestDevice device : context.getDevices()) {
-                                    deviceStates.put(device, FreeDeviceState.AVAILABLE);
-                                }
-                                // Device is released early but this is not marked properly in
-                                // context
-                                listener.releaseDevices(context, deviceStates);
-                                RunUtil.getDefault().sleep(500);
-                                return null;
+        mFakeDeviceManager.setNumDevices(1);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        setCreateConfigExpectations(args);
+        reset(mMockInvocation);
+        when(mMockInvocation.getExitInfo()).thenReturn(null);
+        doAnswer(
+                        invocation -> {
+                            IInvocationContext context =
+                                    (IInvocationContext) invocation.getArguments()[0];
+                            IScheduledInvocationListener listener =
+                                    (IScheduledInvocationListener) invocation.getArguments()[3];
+                            Map<ITestDevice, FreeDeviceState> deviceStates = new HashMap<>();
+                            for (ITestDevice device : context.getDevices()) {
+                                deviceStates.put(device, FreeDeviceState.AVAILABLE);
                             }
-                        });
-        mMockConfiguration.validateOptions();
-        EasyMock.expectLastCall().times(2);
-        replayMocks();
+                            // Device is released early but this is not marked properly in
+                            // context
+                            listener.releaseDevices(context, deviceStates);
+                            RunUtil.getDefault().sleep(500);
+                            return null;
+                        })
+                .when(mMockInvocation)
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
         mScheduler.start();
         mScheduler.addCommand(args);
         RunUtil.getDefault().sleep(100);
@@ -787,8 +808,13 @@
         RunUtil.getDefault().sleep(200);
         mScheduler.shutdown();
         mScheduler.join();
-        verifyMocks();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockConfiguration, times(2)).validateOptions();
+
+        verify(mMockConfigFactory, times(2))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
         assertNotNull(mScheduler.getLastInvocationThrowable());
         assertEquals(
                 "Attempting invocation on device serial0 when one is already running",
@@ -802,18 +828,23 @@
     @Test
     public void testDeviceReleased_unavailable() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesCustom(1, TestDeviceState.NOT_AVAILABLE, IDevice.class);
-        assertEquals(1, mMockManager.getQueueOfAvailableDeviceSize());
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesCustom(1, TestDeviceState.NOT_AVAILABLE, IDevice.class);
+        assertEquals(1, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 0);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 0);
     }
 
     /**
@@ -822,54 +853,69 @@
     @Test
     public void testDeviceReleased_unavailableMulti() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesCustom(2, TestDeviceState.NOT_AVAILABLE, IDevice.class);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 2);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesCustom(2, TestDeviceState.NOT_AVAILABLE, IDevice.class);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 2);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
     }
 
     /** Test that the TCP device NOT available are NOT released. */
     @Test
     public void testTcpDevice_NotReleased() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesStub(1, TestDeviceState.NOT_AVAILABLE, new TcpDevice("serial"));
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesStub(
+                1, TestDeviceState.NOT_AVAILABLE, new TcpDevice("serial"));
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
     }
 
     /** Test that the TCP device NOT available selected for a run is NOT released. */
     @Test
     public void testTcpDevice_NotReleasedMulti() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesStub(2, TestDeviceState.NOT_AVAILABLE, new TcpDevice("serial"));
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 2);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesStub(
+                2, TestDeviceState.NOT_AVAILABLE, new TcpDevice("serial"));
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 2);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 2);
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 2);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
     }
 
     /** Test that the Stub device NOT available are NOT released. */
@@ -877,37 +923,46 @@
     public void testStubDevice_NotReleased() throws Throwable {
         String[] args = new String[] {"test"};
         IDevice stub = new StubDevice("emulator-5554", true);
-        mMockManager.setNumDevicesStub(1, TestDeviceState.NOT_AVAILABLE, stub);
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesStub(1, TestDeviceState.NOT_AVAILABLE, stub);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        assertTrue(mMockManager.getQueueOfAvailableDeviceSize() == 1);
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
+        assertTrue(mFakeDeviceManager.getQueueOfAvailableDeviceSize() == 1);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
     }
 
     /** Test that a device recovery state is reset when returned to the available queue. */
     @Test
     public void testDeviceRecoveryState() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesCustomRealNoRecovery(1, IDevice.class);
-        assertEquals(1, mMockManager.getQueueOfAvailableDeviceSize());
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesCustomRealNoRecovery(1, IDevice.class);
+        assertEquals(1, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
-        assertEquals(1, mMockManager.getQueueOfAvailableDeviceSize());
-        ITestDevice t = mMockManager.allocateDevice();
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
+        assertEquals(1, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
+        ITestDevice t = mFakeDeviceManager.allocateDevice();
         assertTrue(t.getRecoveryMode().equals(RecoveryMode.AVAILABLE));
     }
 
@@ -915,19 +970,24 @@
     @Test
     public void testDevice_unresponsive() throws Throwable {
         String[] args = new String[] {"test"};
-        mMockManager.setNumDevicesUnresponsive(1);
-        assertEquals(1, mMockManager.getQueueOfAvailableDeviceSize());
-        setCreateConfigExpectations(args, 1);
-        setExpectedInvokeCalls(1);
+        mFakeDeviceManager.setNumDevicesUnresponsive(1);
+        assertEquals(1, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join();
-        EasyMock.verify(mMockConfigFactory, mMockConfiguration, mMockInvocation);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any());
+
         // Device does not return to the list since it's unavailable.
-        assertEquals(0, mMockManager.getQueueOfAvailableDeviceSize());
+        assertEquals(0, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
     }
 
     /**
@@ -937,19 +997,27 @@
     @Test
     public void testDisplayCommandQueue() throws Throwable {
         String[] args = new String[] {"empty"};
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        replayMocks();
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
         OutputStream res = new ByteArrayOutputStream();
         PrintWriter pw = new PrintWriter(res);
         mScheduler.displayCommandQueue(pw);
-        verifyMocks();
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
         pw.flush();
-        assertEquals("Id  Config  Created  Exec time  State            Sleep time  Rescheduled  "
-                + "Loop   \n1   empty   0m:00    0m:00      Wait_for_device  N/A         false  "
-                + "      false  \n", res.toString());
+        assertEquals(
+                "Id  Config  Created  Exec time  State            Sleep time  Rescheduled  Loop  "
+                    + " \n"
+                    + "1   empty   0m:00    0m:00      Wait_for_device  N/A         false       "
+                    + " false  \n",
+                res.toString());
         mScheduler.shutdown();
     }
 
@@ -962,14 +1030,20 @@
         String[] args = new String[] {"empty"};
         OutputStream res = new ByteArrayOutputStream();
         PrintWriter pw = new PrintWriter(res);
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        mMockConfiguration.dumpXml(EasyMock.anyObject());
-        replayMocks();
+        setCreateConfigExpectations(args);
+
+        mMockConfiguration.dumpXml(any());
+
         mScheduler.start();
         mScheduler.addCommand(args);
         mScheduler.dumpCommandsXml(pw, null);
-        verifyMocks();
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
         pw.flush();
         String filename = res.toString().replace("Saved command dump to ", "").trim();
         File test = new File(filename);
@@ -988,15 +1062,20 @@
     @Test
     public void testDisplayCommandsInfo() throws Throwable {
         String[] args = new String[] {"empty"};
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        replayMocks();
+        setCreateConfigExpectations(args);
+
         mScheduler.start();
         mScheduler.addCommand(args);
         OutputStream res = new ByteArrayOutputStream();
         PrintWriter pw = new PrintWriter(res);
         mScheduler.displayCommandsInfo(pw, null);
-        verifyMocks();
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
         pw.flush();
         assertEquals("Command 1: [0m:00] empty\n", res.toString());
         mScheduler.shutdown();
@@ -1009,9 +1088,9 @@
     @Test
     public void testGetInvocationInfo_null() throws Throwable {
         String[] args = new String[] {"empty", "test"};
-        setCreateConfigExpectations(args, 1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         mScheduler.addCommand(args);
         assertNull(mScheduler.getInvocationInfo(999));
@@ -1021,13 +1100,13 @@
     @Test
     public void testAllocateDevices() throws Exception {
         String[] args = new String[] {"foo", "test"};
-        mMockManager.setNumDevices(1);
-        setCreateConfigExpectations(args, 1);
+        mFakeDeviceManager.setNumDevices(1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         DeviceAllocationResult results =
-                mScheduler.allocateDevices(mMockConfiguration, mMockManager);
+                mScheduler.allocateDevices(mMockConfiguration, mFakeDeviceManager);
         assertTrue(results.wasAllocationSuccessful());
         Map<String, ITestDevice> devices = results.getAllocatedDevices();
         assertEquals(1, devices.size());
@@ -1037,8 +1116,8 @@
     @Test
     public void testAllocateDevices_replicated() throws Exception {
         String[] args = new String[] {"foo", "test"};
-        mMockManager.setNumDevices(3);
-        setCreateConfigExpectations(args, 1);
+        mFakeDeviceManager.setNumDevices(3);
+        setCreateConfigExpectations(args);
         OptionSetter setter = new OptionSetter(mCommandOptions);
         setter.setOptionValue("replicate-parent-setup", "true");
         mCommandOptions.setShardCount(3);
@@ -1046,16 +1125,15 @@
         for (int i = 0; i < 2; i++) {
             IConfiguration configReplicat = new Configuration("test", "test");
             configReplicat.setDeviceConfig(new DeviceConfigurationHolder("serial"));
-            EasyMock.expect(
-                            mMockConfiguration.partialDeepClone(
-                                    Arrays.asList(Configuration.DEVICE_NAME), mMockKeyStoreClient))
-                    .andReturn(configReplicat);
+            when(mMockConfiguration.partialDeepClone(
+                            Arrays.asList(Configuration.DEVICE_NAME), mMockKeyStoreClient))
+                    .thenReturn(configReplicat);
         }
-        mMockConfiguration.setDeviceConfigList(EasyMock.anyObject());
-        replayMocks();
+        mMockConfiguration.setDeviceConfigList(any());
+
         mScheduler.start();
         DeviceAllocationResult results =
-                mScheduler.allocateDevices(mMockConfiguration, mMockManager);
+                mScheduler.allocateDevices(mMockConfiguration, mFakeDeviceManager);
         assertTrue(results.wasAllocationSuccessful());
         Map<String, ITestDevice> devices = results.getAllocatedDevices();
         // With replicated setup, all devices get allocated.
@@ -1075,20 +1153,20 @@
     public void testAllocateDevices_multipleDevices() throws Exception {
         String[] args = new String[] {"foo", "test"};
 
-        mMockManager.setNumDevices(2);
-        mMockDeviceConfig.add(createDeviceConfig("serial0"));
-        mMockDeviceConfig.add(createDeviceConfig("serial1"));
+        mFakeDeviceManager.setNumDevices(2);
+        mDeviceConfigList.add(createDeviceConfig("serial0"));
+        mDeviceConfigList.add(createDeviceConfig("serial1"));
 
-        setCreateConfigExpectations(args, 1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         DeviceAllocationResult results =
-                mScheduler.allocateDevices(mMockConfiguration, mMockManager);
+                mScheduler.allocateDevices(mMockConfiguration, mFakeDeviceManager);
         assertTrue(results.wasAllocationSuccessful());
         Map<String, ITestDevice> devices = results.getAllocatedDevices();
         assertEquals(2, devices.size());
-        assertEquals(0, mMockManager.getQueueOfAvailableDeviceSize());
+        assertEquals(0, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
         mScheduler.shutdown();
     }
 
@@ -1096,20 +1174,20 @@
     public void testAllocateDevices_multipleDevices_failed() throws Exception {
         String[] args = new String[] {"foo", "test"};
 
-        mMockManager.setNumDevices(2);
-        mMockDeviceConfig.add(createDeviceConfig("serial0"));
-        mMockDeviceConfig.add(createDeviceConfig("not_exist_serial"));
+        mFakeDeviceManager.setNumDevices(2);
+        mDeviceConfigList.add(createDeviceConfig("serial0"));
+        mDeviceConfigList.add(createDeviceConfig("not_exist_serial"));
 
-        setCreateConfigExpectations(args, 1);
+        setCreateConfigExpectations(args);
         mMockConfiguration.validateOptions();
-        replayMocks();
+
         mScheduler.start();
         DeviceAllocationResult results =
-                mScheduler.allocateDevices(mMockConfiguration, mMockManager);
+                mScheduler.allocateDevices(mMockConfiguration, mFakeDeviceManager);
         assertFalse(results.wasAllocationSuccessful());
         Map<String, ITestDevice> devices = results.getAllocatedDevices();
         assertEquals(0, devices.size());
-        assertEquals(2, mMockManager.getQueueOfAvailableDeviceSize());
+        assertEquals(2, mFakeDeviceManager.getQueueOfAvailableDeviceSize());
         mScheduler.shutdown();
     }
 
@@ -1120,31 +1198,43 @@
     @Test
     @SuppressWarnings("unchecked")
     public void testExecCommand_multipleDevices() throws Throwable {
-        String[] args = new String[] {
-            "foo"
-        };
-        mMockManager.setNumDevices(2);
-        mMockDeviceConfig.add(createDeviceConfig("serial0"));
-        mMockDeviceConfig.add(createDeviceConfig("serial1"));
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        mMockInvocation.invoke((IInvocationContext)EasyMock.anyObject(),
-                (IConfiguration)EasyMock.anyObject(), (IRescheduler)EasyMock.anyObject(),
-                (ITestInvocationListener)EasyMock.anyObject(),
+        String[] args = new String[] {"foo"};
+        mFakeDeviceManager.setNumDevices(2);
+        mDeviceConfigList.add(createDeviceConfig("serial0"));
+        mDeviceConfigList.add(createDeviceConfig("serial1"));
+        setCreateConfigExpectations(args);
+
+        mMockInvocation.invoke(
+                (IInvocationContext) any(),
+                (IConfiguration) any(),
+                (IRescheduler) any(),
+                (ITestInvocationListener) any(),
                 // This is FreeDeviceHandler.
-                (IScheduledInvocationListener)EasyMock.anyObject());
-        IScheduledInvocationListener mockListener = EasyMock
-                .createMock(IScheduledInvocationListener.class);
-        mockListener.invocationInitiated((IInvocationContext) EasyMock.anyObject());
-        mockListener.invocationComplete((IInvocationContext)EasyMock.anyObject(),
-                (Map<ITestDevice, FreeDeviceState>)EasyMock.anyObject());
-        replayMocks(mockListener);
+                (IScheduledInvocationListener) any());
+        IScheduledInvocationListener mockListener = mock(IScheduledInvocationListener.class);
+        mockListener.invocationInitiated((IInvocationContext) any());
+        mockListener.invocationComplete(
+                (IInvocationContext) any(), (Map<ITestDevice, FreeDeviceState>) any());
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any(),
+                        // This is FreeDeviceHandler.
+                        (IScheduledInvocationListener) any());
 
         mScheduler.start();
         mScheduler.execCommand(mockListener, args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join(2 * 1000);
-        verifyMocks(mockListener);
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -1153,17 +1243,13 @@
      */
     @Test
     public void testExecCommand_multipleDevices_noDevice() throws Throwable {
-        String[] args = new String[] {
-            "foo"
-        };
-        mMockManager.setNumDevices(2);
-        mMockDeviceConfig.add(createDeviceConfig("serial0"));
-        mMockDeviceConfig.add(createDeviceConfig("not_exist_serial"));
-        setCreateConfigExpectations(args, 1);
-        mMockConfiguration.validateOptions();
-        IScheduledInvocationListener mockListener = EasyMock
-                .createMock(IScheduledInvocationListener.class);
-        replayMocks(mockListener);
+        String[] args = new String[] {"foo"};
+        mFakeDeviceManager.setNumDevices(2);
+        mDeviceConfigList.add(createDeviceConfig("serial0"));
+        mDeviceConfigList.add(createDeviceConfig("not_exist_serial"));
+        setCreateConfigExpectations(args);
+
+        IScheduledInvocationListener mockListener = mock(IScheduledInvocationListener.class);
 
         mScheduler.start();
         try {
@@ -1174,7 +1260,13 @@
         }
         mScheduler.shutdownOnEmpty();
         mScheduler.join(2 * 1000);
-        verifyMocks(mockListener);
+
+        verify(mMockConfiguration).validateOptions();
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
     }
 
     /**
@@ -1187,34 +1279,24 @@
                 new String[] {
                     "foo", "--invocation-data", "test",
                 };
-        setCreateConfigExpectations(args, 1);
+        setCreateConfigExpectations(args);
         OptionSetter setter = new OptionSetter(mCommandOptions);
         // If invocation-data are added and we are in a versioned invocation, the data should not
         // be added again.
         setter.setOptionValue("invocation-data", "key", "value");
-        mMockConfigDescriptor.setSandboxed(true);
-
-        mMockInvocation.invoke(
-                (IInvocationContext) EasyMock.anyObject(),
-                (IConfiguration) EasyMock.anyObject(),
-                (IRescheduler) EasyMock.anyObject(),
-                (ITestInvocationListener) EasyMock.anyObject(),
-                EasyMock.anyObject());
-        EasyMock.expectLastCall().times(1);
+        mConfigDescriptor.setSandboxed(true);
 
         mMockConfiguration.validateOptions();
-        IDevice mockIDevice = EasyMock.createMock(IDevice.class);
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        EasyMock.expect(mockDevice.getSerialNumber()).andStubReturn("serial");
-        EasyMock.expect(mockDevice.getDeviceState()).andStubReturn(TestDeviceState.ONLINE);
-        mockDevice.setRecoveryMode(EasyMock.eq(RecoveryMode.AVAILABLE));
-        EasyMock.expect(mockDevice.getIDevice()).andStubReturn(mockIDevice);
-        IScheduledInvocationListener mockListener =
-                EasyMock.createMock(IScheduledInvocationListener.class);
-        mockListener.invocationInitiated((InvocationContext) EasyMock.anyObject());
-        mockListener.invocationComplete(
-                (IInvocationContext) EasyMock.anyObject(), EasyMock.anyObject());
-        EasyMock.expect(mockDevice.waitForDeviceShell(EasyMock.anyLong())).andReturn(true);
+        IDevice mockIDevice = mock(IDevice.class);
+        ITestDevice mockDevice = mock(ITestDevice.class);
+        when(mockDevice.getSerialNumber()).thenReturn("serial");
+        when(mockDevice.getDeviceState()).thenReturn(TestDeviceState.ONLINE);
+        mockDevice.setRecoveryMode(eq(RecoveryMode.AVAILABLE));
+        when(mockDevice.getIDevice()).thenReturn(mockIDevice);
+        IScheduledInvocationListener mockListener = mock(IScheduledInvocationListener.class);
+        mockListener.invocationInitiated((InvocationContext) any());
+        mockListener.invocationComplete((IInvocationContext) any(), any());
+        when(mockDevice.waitForDeviceShell(anyLong())).thenReturn(true);
 
         mScheduler =
                 new TestableCommandScheduler() {
@@ -1230,12 +1312,23 @@
                     }
                 };
 
-        replayMocks(mockDevice, mockListener);
         mScheduler.start();
         mScheduler.execCommand(mockListener, args);
         mScheduler.shutdownOnEmpty();
         mScheduler.join(2 * 1000);
-        verifyMocks(mockListener);
+        verify(mMockInvocation, times(1))
+                .invoke(
+                        (IInvocationContext) any(),
+                        (IConfiguration) any(),
+                        (IRescheduler) any(),
+                        (ITestInvocationListener) any(),
+                        // This is FreeDeviceHandler.
+                        (IScheduledInvocationListener) any());
+
+        verify(mMockConfigFactory, times(1))
+                .createConfigurationFromArgs(
+                        AdditionalMatchers.aryEq(args), isNull(), (IKeyStoreClient) any());
+        mFakeDeviceManager.assertDevicesFreed();
 
         // only attribute is invocation ID
         assertEquals(1, mContext.getAttributes().size());
diff --git a/javatests/com/android/tradefed/device/TestDeviceTest.java b/javatests/com/android/tradefed/device/TestDeviceTest.java
index b7b7ccd..f46166e 100644
--- a/javatests/com/android/tradefed/device/TestDeviceTest.java
+++ b/javatests/com/android/tradefed/device/TestDeviceTest.java
@@ -5019,6 +5019,40 @@
         }
     }
 
+    public void testGetFoldableStates() throws Exception {
+        mTestDevice = new TestableTestDevice() {
+            @Override
+            public CommandResult executeShellV2Command(String cmd)
+                    throws DeviceNotAvailableException {
+                CommandResult result = new CommandResult(CommandStatus.SUCCESS);
+                result.setStdout("Supported states: [\n" +
+                        " DeviceState{identifier=0, name='CLOSED'},\n" +
+                        " DeviceState{identifier=1, name='HALF_OPENED'},\n" +
+                        " DeviceState{identifier=2, name='OPENED'},\n" +
+                        "]\n");
+                return result;
+            }
+        };
+
+        Set<DeviceFoldableState> states = mTestDevice.getFoldableStates();
+        assertEquals(3, states.size());
+    }
+
+    public void testGetCurrentFoldableState() throws Exception {
+        mTestDevice = new TestableTestDevice() {
+            @Override
+            public CommandResult executeShellV2Command(String cmd)
+                    throws DeviceNotAvailableException {
+                CommandResult result = new CommandResult(CommandStatus.SUCCESS);
+                result.setStdout("Committed state: DeviceState{identifier=2, name='DEFAULT'}\n");
+                return result;
+            }
+        };
+
+        DeviceFoldableState state = mTestDevice.getCurrentFoldableState();
+        assertEquals(2, state.getIdentifier());
+    }
+
     private IExpectationSetters<CommandResult> setGetPropertyExpectation(
             String property, String value) {
         CommandResult stubResult = new CommandResult(CommandStatus.SUCCESS);
diff --git a/javatests/com/android/tradefed/testtype/rust/RustBinaryHostTestTest.java b/javatests/com/android/tradefed/testtype/rust/RustBinaryHostTestTest.java
index d72b25e..c826faf 100644
--- a/javatests/com/android/tradefed/testtype/rust/RustBinaryHostTestTest.java
+++ b/javatests/com/android/tradefed/testtype/rust/RustBinaryHostTestTest.java
@@ -159,6 +159,7 @@
     }
 
     private void mockTestRunExpect(File binary, CommandResult res) throws Exception {
+        mMockRunUtil.setWorkingDir(binary.getParentFile());
         EasyMock.expect(
                         mMockRunUtil.runTimedCmd(
                                 EasyMock.anyLong(), EasyMock.eq(binary.getAbsolutePath())))
@@ -167,6 +168,7 @@
 
     private void mockBenchmarkRunExpect(File binary, String output) throws Exception {
         CommandResult res = newCommandResult(CommandStatus.SUCCESS, "", "");
+        mMockRunUtil.setWorkingDir(binary.getParentFile());
         EasyMock.expect(
                         mMockRunUtil.runTimedCmd(
                                 EasyMock.anyLong(),
@@ -328,6 +330,7 @@
             mockListenerStarted(binary, 9);
             mockListenerLog(binary, false);
             CommandResult res = successResult("", resultCount(6, 1, 2));
+            mMockRunUtil.setWorkingDir(binary.getParentFile());
             EasyMock.expect(
                             mMockRunUtil.runTimedCmd(
                                     EasyMock.anyLong(),
@@ -372,6 +375,7 @@
 
             mockListenerLog(binary, false);
             CommandResult res = successResult("", resultCount(3, 0, 0));
+            mMockRunUtil.setWorkingDir(binary.getParentFile());
             EasyMock.expect(
                             mMockRunUtil.runTimedCmd(
                                     EasyMock.anyLong(),
@@ -433,6 +437,7 @@
             // Multiple include filters are run one by one.
             mockListenerLog(binary, false);
             CommandResult res = successResult("", resultCount(2, 0, 0));
+            mMockRunUtil.setWorkingDir(binary.getParentFile());
             EasyMock.expect(
                             mMockRunUtil.runTimedCmd(
                                     EasyMock.anyLong(),
@@ -446,6 +451,7 @@
             mMockListener.testRunFailed("Test run incomplete. Started 2 tests, finished 0");
             mockListenerLog(binary, false);
             res = successResult("", resultCount(3, 0, 0));
+            mMockRunUtil.setWorkingDir(binary.getParentFile());
             EasyMock.expect(
                             mMockRunUtil.runTimedCmd(
                                     EasyMock.anyLong(),
diff --git a/javatests/com/android/tradefed/testtype/suite/AtestRunnerTest.java b/javatests/com/android/tradefed/testtype/suite/AtestRunnerTest.java
index c91e3f0..051e228 100644
--- a/javatests/com/android/tradefed/testtype/suite/AtestRunnerTest.java
+++ b/javatests/com/android/tradefed/testtype/suite/AtestRunnerTest.java
@@ -20,6 +20,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
 
 import com.android.tradefed.build.DeviceBuildInfo;
 import com.android.tradefed.build.IDeviceBuildInfo;
@@ -34,12 +35,15 @@
 import com.android.tradefed.testtype.InstrumentationTest;
 import com.android.tradefed.testtype.UiAutomatorTest;
 import com.android.tradefed.util.AbiUtils;
+import com.android.tradefed.util.CommandResult;
+import com.android.tradefed.util.CommandStatus;
 import com.android.tradefed.util.FileUtil;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -85,6 +89,13 @@
         abis.add(new Abi(ABI, AbiUtils.getBitness(ABI)));
         doReturn(abis).when(mSpyRunner).getAbis(mMockDevice);
         doReturn(new File("some-dir")).when(mBuildInfo).getTestsDir();
+
+        CommandResult result = new CommandResult(CommandStatus.SUCCESS);
+        result.setStdout("Supported states: [\n" +
+                " DeviceState{identifier=0, name='DEFAULT'},\n" +
+                "]\n");
+        when(mMockDevice.executeShellV2Command("cmd device_state print-states"))
+                .thenReturn(result);
     }
 
     @Test
diff --git a/javatests/com/android/tradefed/testtype/suite/BaseTestSuiteTest.java b/javatests/com/android/tradefed/testtype/suite/BaseTestSuiteTest.java
index ee14ac1..1978739 100644
--- a/javatests/com/android/tradefed/testtype/suite/BaseTestSuiteTest.java
+++ b/javatests/com/android/tradefed/testtype/suite/BaseTestSuiteTest.java
@@ -24,6 +24,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.android.ddmlib.IDevice;
 import com.android.tradefed.build.DeviceBuildInfo;
 import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.config.ConfigurationDef;
@@ -84,6 +85,8 @@
 
         EasyMock.expect(mMockDevice.getProperty(EasyMock.anyObject())).andReturn("arm64-v8a");
         EasyMock.expect(mMockDevice.getProperty(EasyMock.anyObject())).andReturn("armeabi-v7a");
+        EasyMock.expect(mMockDevice.getIDevice()).andStubReturn(EasyMock.createMock(IDevice.class));
+        EasyMock.expect(mMockDevice.getFoldableStates()).andStubReturn(new HashSet<>());
         EasyMock.replay(mMockDevice);
     }
 
@@ -394,17 +397,14 @@
      */
     @Test
     public void testLoadTestsForMultiAbi() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-abi");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(2, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stubAbi"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/stubAbi"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -413,8 +413,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters");
@@ -423,14 +421,13 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         // We only create the primary abi of the parameterized module version.
         assertEquals(3, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized"));
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized[instant]"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/stub-parameterized"));
-        EasyMock.verify(mockDevice);
 
         TestSuiteStub testSuiteStub =
                 (TestSuiteStub)
@@ -455,8 +452,6 @@
     /** Ensure parameterized modules are created properly even when main abi is filtered. */
     @Test
     public void testLoadTests_parameterizedModule_load_with_filter() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         Set<String> excludeModule = new HashSet<>();
         excludeModule.add("arm64-v8a suite/load-filter-test");
         mRunner.setExcludeFilter(excludeModule);
@@ -468,21 +463,18 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(2, configMap.size());
         // Config main abi non-parameterized is filtered, this shouldn't prevent the parameterized
         // version from being created, and the other abi.
         assertTrue(configMap.containsKey("arm64-v8a suite/load-filter-test[instant]"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/load-filter-test"));
-        EasyMock.verify(mockDevice);
     }
 
     /** Ensure parameterized modules are filtered when requested. */
     @Test
     public void testLoadTests_parameterizedModule_load_with_filter_param() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         Set<String> excludeModule = new HashSet<>();
         excludeModule.add("arm64-v8a suite/load-filter-test[instant]");
         mRunner.setExcludeFilter(excludeModule);
@@ -494,14 +486,13 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(2, configMap.size());
         // Config main abi parameterized is filtered, this shouldn't prevent the non-parameterized
         // version from being created, and the other abi.
         assertTrue(configMap.containsKey("arm64-v8a suite/load-filter-test"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/load-filter-test"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -510,8 +501,6 @@
      */
     @Test
     public void testLoadTests_forcedModule_load_with_filter_param() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         mRunner.setModuleParameter(ModuleParameters.INSTANT_APP);
         Set<String> includeModule = new HashSet<>();
         includeModule.add("arm64-v8a suite/load-filter-test");
@@ -520,13 +509,12 @@
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "test-filter-load");
         setter.setOptionValue("enable-parameterized-modules", "true");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         // Config main abi parameterized is filtered, this shouldn't prevent the non-parameterized
         // version from being created, and the other abi.
         assertTrue(configMap.containsKey("arm64-v8a suite/load-filter-test[instant]"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -535,8 +523,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_multiAbi() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-abi");
@@ -545,7 +531,7 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(5, configMap.size());
         // stub-parameterized-abi2 is not parameterized so by default both abi are created.
@@ -555,13 +541,10 @@
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi"));
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi[instant]"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/stub-parameterized-abi"));
-        EasyMock.verify(mockDevice);
     }
 
     @Test
     public void testLoadTests_parameterizedModule_multiAbi_filter() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         Set<String> includeFilters = new HashSet<>();
         includeFilters.add("suite/stub-parameterized-abi[instant]");
         mRunner.setIncludeFilter(includeFilters);
@@ -573,11 +556,10 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi[instant]"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -586,8 +568,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_multiAbi_forced() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-abi");
@@ -597,17 +577,14 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi[instant]"));
-        EasyMock.verify(mockDevice);
     }
 
     @Test
     public void testLoadTests_parameterizedModule_only_instant() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-abi-alone");
@@ -617,10 +594,9 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(0, configMap.size());
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -629,8 +605,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_multiAbi_forcedNotInstant() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-abi");
@@ -640,7 +614,7 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(4, configMap.size());
         // stub-parameterized-abi2 is not parameterized so by default both abi are created.
@@ -650,7 +624,6 @@
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi"));
         //assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi[instant]"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/stub-parameterized-abi"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -659,8 +632,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_notMultiAbi() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-not-multi");
@@ -669,19 +640,15 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(2, configMap.size());
         // stub-parameterized-abi is parameterized and not multi_abi so it creates only one abi
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi4"));
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi4[instant]"));
-        EasyMock.verify(mockDevice);
     }
 
     @Test
     public void testLoadTests_parameterizedModule_notMultiAbi_withFilter() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-not-multi");
@@ -693,18 +660,15 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         // stub-parameterized-abi is parameterized and not multi_abi so it creates only one abi
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi4"));
-        EasyMock.verify(mockDevice);
     }
 
     @Test
     public void testLoadTests_parameterizedModule_filter() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-not-multi");
@@ -714,13 +678,12 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         // stub-parameterized-abi is parameterized and not multi_abi so it creates only one abi
         // instant_app is filtered so only the regular version of it is created.
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi4"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -729,8 +692,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_forced() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-not-multi");
@@ -740,12 +701,11 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(1, configMap.size());
         // stub-parameterized-abi is parameterized and not multi_abi so it creates only one abi
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized-abi4[instant]"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -754,13 +714,11 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_mutuallyExclusiveFamily() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters-fail");
         setter.setOptionValue("enable-parameterized-modules", "true");
-        EasyMock.replay(mockDevice);
+
         try {
             mRunner.loadTests();
             fail("Should have thrown an exception.");
@@ -772,14 +730,11 @@
                             + "not_instant_app and instant_app when only one expected.'",
                     expected.getMessage());
         }
-        EasyMock.verify(mockDevice);
     }
 
     /** Test that loading the option parameterization is gated by the option. */
     @Test
     public void testLoadTests_optionalParameterizedModule() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters");
@@ -789,21 +744,18 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         assertEquals(4, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized"));
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized[instant]"));
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized[secondary_user]"));
         assertTrue(configMap.containsKey("armeabi-v7a suite/stub-parameterized"));
-        EasyMock.verify(mockDevice);
     }
 
     /** Test that we can explicitly request the option parameterization type. */
     @Test
     public void testLoadTests_optionalParameterizedModule_filter() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters");
@@ -814,12 +766,11 @@
                 "test-arg",
                 "com.android.tradefed.testtype.suite.TestSuiteStub:"
                         + "exclude-annotation:android.platform.test.annotations.AppModeInstant");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         // Only the secondary_user requested is created
         assertEquals(1, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized[secondary_user]"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
@@ -828,8 +779,6 @@
      */
     @Test
     public void testLoadTests_parameterizedModule_optionSetupAfterConfigAdded() throws Exception {
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
         OptionSetter setter = new OptionSetter(mRunner);
         setter.setOptionValue("suite-config-prefix", "suite");
         setter.setOptionValue("run-suite-tag", "example-suite-parameters");
@@ -839,12 +788,11 @@
         setter.setOptionValue(
                 "test-arg",
                 "com.android.tradefed.targetprep.CreateUserPreparer:reuse-test-user:true");
-        EasyMock.replay(mockDevice);
+
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
         // We only create the primary abi of the parameterized module version.
         assertEquals(1, configMap.size());
         assertTrue(configMap.containsKey("arm64-v8a suite/stub-parameterized[secondary_user]"));
-        EasyMock.verify(mockDevice);
 
         IConfiguration config = configMap.get("arm64-v8a suite/stub-parameterized[secondary_user]");
 
diff --git a/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java b/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
index f14543a..7bf1f70 100644
--- a/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
+++ b/javatests/com/android/tradefed/testtype/suite/TestMappingSuiteRunnerTest.java
@@ -19,6 +19,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
+import com.android.ddmlib.IDevice;
 import com.android.tradefed.build.BuildInfoKey.BuildInfoFileKey;
 import com.android.tradefed.build.IDeviceBuildInfo;
 import com.android.tradefed.config.Configuration;
@@ -133,6 +134,8 @@
         EasyMock.expect(mBuildInfo.getTestsDir()).andReturn(new File(NON_EXISTING_DIR)).anyTimes();
         EasyMock.expect(mMockDevice.getProperty(EasyMock.anyObject())).andReturn(ABI_1);
         EasyMock.expect(mMockDevice.getProperty(EasyMock.anyObject())).andReturn(ABI_2);
+        EasyMock.expect(mMockDevice.getIDevice()).andStubReturn(EasyMock.createMock(IDevice.class));
+        EasyMock.expect(mMockDevice.getFoldableStates()).andStubReturn(new HashSet<>());
         EasyMock.replay(mBuildInfo, mMockDevice);
     }
 
@@ -671,16 +674,10 @@
     public void testLoadTestsForMultiAbi() throws Exception {
         mOptionSetter.setOptionValue("include-filter", "suite/stubAbi");
 
-        ITestDevice mockDevice = EasyMock.createMock(ITestDevice.class);
-        mRunner.setDevice(mockDevice);
-        EasyMock.replay(mockDevice);
-
         LinkedHashMap<String, IConfiguration> configMap = mRunner.loadTests();
-
         assertEquals(2, configMap.size());
         assertTrue(configMap.containsKey(ABI_1 + " suite/stubAbi"));
         assertTrue(configMap.containsKey(ABI_2 + " suite/stubAbi"));
-        EasyMock.verify(mockDevice);
     }
 
     /**
diff --git a/javatests/com/android/tradefed/testtype/suite/params/ModuleParametersHelperTest.java b/javatests/com/android/tradefed/testtype/suite/params/ModuleParametersHelperTest.java
index 0ef4671..472f481 100644
--- a/javatests/com/android/tradefed/testtype/suite/params/ModuleParametersHelperTest.java
+++ b/javatests/com/android/tradefed/testtype/suite/params/ModuleParametersHelperTest.java
@@ -82,7 +82,7 @@
     public void resolveParamString_notGroupParam_returnsSetOfSameParam() {
         Map<ModuleParameters, IModuleParameterHandler> resolvedParams =
                 ModuleParametersHelper.resolveParam(
-                        ModuleParameters.INSTANT_APP.toString(), /* withOptional= */ true);
+                        ModuleParameters.INSTANT_APP, /* withOptional= */ true);
 
         assertEquals(resolvedParams.size(), 1);
         assertEquals(resolvedParams.keySet().iterator().next(), ModuleParameters.INSTANT_APP);
@@ -92,7 +92,7 @@
     public void resolveParamString_groupParam_returnsSetOfMultipleParams() {
         Map<ModuleParameters, IModuleParameterHandler> resolvedParams =
                 ModuleParametersHelper.resolveParam(
-                        ModuleParameters.MULTIUSER.toString(), /* withOptional= */ true);
+                        ModuleParameters.MULTIUSER, /* withOptional= */ true);
 
         assertNotEquals(resolvedParams.size(), 1);
         assertFalse(resolvedParams.keySet().contains(ModuleParameters.MULTIUSER));
diff --git a/src/com/android/tradefed/config/ConfigurationFactory.java b/src/com/android/tradefed/config/ConfigurationFactory.java
index aac2139..921918c 100644
--- a/src/com/android/tradefed/config/ConfigurationFactory.java
+++ b/src/com/android/tradefed/config/ConfigurationFactory.java
@@ -605,6 +605,19 @@
 
         Map<String, String> uniqueMap =
                 extractTemplates(configName, listArgs, optionArgsRef, keyStoreClient);
+        if (allowedObjects != null && !allowedObjects.isEmpty()) {
+            ConfigLoader tmpLoader = new ConfigLoader(false);
+            // For partial loading be lenient about templates and let the delegate deal with it.
+            // In some cases this won't be 100% correct but it's better than failing on all new
+            // configs.
+            for (String key : uniqueMap.keySet()) {
+                try {
+                    tmpLoader.findConfigName(uniqueMap.get(key), null);
+                } catch (ConfigurationException e) {
+                    uniqueMap.put(key, "empty");
+                }
+            }
+        }
         ConfigurationDef configDef = getConfigurationDef(configName, false, uniqueMap);
         if (!uniqueMap.isEmpty()) {
             // remove the bad ConfigDef from the cache.
diff --git a/src/com/android/tradefed/device/NativeDevice.java b/src/com/android/tradefed/device/NativeDevice.java
index 49bec5f..5538bb0 100644
--- a/src/com/android/tradefed/device/NativeDevice.java
+++ b/src/com/android/tradefed/device/NativeDevice.java
@@ -5268,6 +5268,16 @@
         return tombstones;
     }
 
+    @Override
+    public Set<DeviceFoldableState> getFoldableStates() throws DeviceNotAvailableException {
+        throw new UnsupportedOperationException("No support for foldable states.");
+    }
+
+    @Override
+    public DeviceFoldableState getCurrentFoldableState() throws DeviceNotAvailableException {
+        throw new UnsupportedOperationException("No support for foldable states.");
+    }
+
     /** Validate that pid is an integer and not empty. */
     private boolean checkValidPid(String output) {
         if (output.isEmpty()) {
diff --git a/src/com/android/tradefed/device/TestDevice.java b/src/com/android/tradefed/device/TestDevice.java
index 5c4ab05..bb6808c 100644
--- a/src/com/android/tradefed/device/TestDevice.java
+++ b/src/com/android/tradefed/device/TestDevice.java
@@ -47,6 +47,7 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -1951,4 +1952,46 @@
 
         return displays;
     }
+
+    @Override
+    public Set<DeviceFoldableState> getFoldableStates() throws DeviceNotAvailableException {
+        if (getIDevice() instanceof StubDevice) {
+            return new HashSet<>();
+        }
+        CommandResult result = executeShellV2Command("cmd device_state print-states");
+        if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
+            // Can't throw an exception since it would fail on non-supported version
+            CLog.w("Failed to enumerate foldable configurations. stderr: %s", result.getStderr());
+            return new HashSet<>();
+        }
+        Set<DeviceFoldableState> foldableStates = new LinkedHashSet<>();
+        Pattern deviceStatePattern =
+                Pattern.compile("DeviceState\\{identifier=(\\d+), name='(\\S+)'\\}\\S*");
+        for (String line : result.getStdout().split("\n")) {
+            Matcher m = deviceStatePattern.matcher(line.trim());
+            if (m.matches()) {
+                foldableStates.add(
+                        new DeviceFoldableState(Integer.parseInt(m.group(1)), m.group(2)));
+            }
+        }
+        return foldableStates;
+    }
+
+    @Override
+    public DeviceFoldableState getCurrentFoldableState() throws DeviceNotAvailableException {
+        if (getIDevice() instanceof StubDevice) {
+            return null;
+        }
+        CommandResult result = executeShellV2Command("cmd device_state state");
+        Pattern deviceStatePattern =
+                Pattern.compile(
+                        "Committed state: DeviceState\\{identifier=(\\d+), name='(\\S+)'\\}\\S*");
+        for (String line : result.getStdout().split("\n")) {
+            Matcher m = deviceStatePattern.matcher(line.trim());
+            if (m.matches()) {
+                return new DeviceFoldableState(Integer.parseInt(m.group(1)), m.group(2));
+            }
+        }
+        return null;
+    }
 }
diff --git a/src/com/android/tradefed/targetprep/suite/FoldableModePreparer.java b/src/com/android/tradefed/targetprep/suite/FoldableModePreparer.java
index 6834925..0a5da38 100644
--- a/src/com/android/tradefed/targetprep/suite/FoldableModePreparer.java
+++ b/src/com/android/tradefed/targetprep/suite/FoldableModePreparer.java
@@ -16,10 +16,12 @@
 package com.android.tradefed.targetprep.suite;
 
 import com.android.tradefed.config.Option;
+import com.android.tradefed.device.DeviceFoldableState;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.error.HarnessRuntimeException;
 import com.android.tradefed.invoker.TestInformation;
 import com.android.tradefed.result.error.DeviceErrorIdentifier;
+import com.android.tradefed.result.error.TestErrorIdentifier;
 import com.android.tradefed.targetprep.BaseTargetPreparer;
 import com.android.tradefed.targetprep.BuildError;
 import com.android.tradefed.targetprep.TargetSetupError;
@@ -31,9 +33,14 @@
  */
 public class FoldableModePreparer extends BaseTargetPreparer {
 
-    @Option(name = "foldable-state-identifier", description = "The integer state identifier of the foldable mode.")
+    @Option(name = "foldable-state-identifier",
+            description = "The integer state identifier of the foldable mode.")
     private Long mStateIdentifier = null;
 
+    @Option(name = "enforce-state",
+            description = "Ensure the state wasn't changed between setup & teardown.")
+    private boolean mValidateSameState = true;
+
     public FoldableModePreparer() {}
 
     public FoldableModePreparer(long stateIdentifier) {
@@ -63,6 +70,16 @@
         if (mStateIdentifier == null) {
             return;
         }
+        if (mValidateSameState) {
+            DeviceFoldableState state = testInformation.getDevice().getCurrentFoldableState();
+            if (state == null || state.getIdentifier() != mStateIdentifier) {
+                throw new HarnessRuntimeException(
+                        String.format("Device foldable state changed to '%s', "
+                                + "it should have been '%s'", state, mStateIdentifier),
+                            TestErrorIdentifier.MODIFIED_FOLDABLE_STATE);
+            }
+        }
+
         CommandResult result = testInformation.getDevice().executeShellV2Command(
                 "cmd device_state state reset");
         if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
diff --git a/src/com/android/tradefed/testtype/suite/BaseTestSuite.java b/src/com/android/tradefed/testtype/suite/BaseTestSuite.java
index c4d1da9..184287a 100644
--- a/src/com/android/tradefed/testtype/suite/BaseTestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/BaseTestSuite.java
@@ -21,7 +21,10 @@
 import com.android.tradefed.config.Option;
 import com.android.tradefed.config.Option.Importance;
 import com.android.tradefed.config.OptionClass;
+import com.android.tradefed.device.DeviceFoldableState;
 import com.android.tradefed.device.DeviceNotAvailableException;
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.device.StubDevice;
 import com.android.tradefed.error.HarnessRuntimeException;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.FileInputStreamSource;
@@ -31,6 +34,7 @@
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.ITestFileFilterReceiver;
 import com.android.tradefed.testtype.ITestFilterReceiver;
+import com.android.tradefed.testtype.suite.params.FoldableExpandingHandler;
 import com.android.tradefed.testtype.suite.params.IModuleParameterHandler;
 import com.android.tradefed.testtype.suite.params.ModuleParameters;
 import com.android.tradefed.testtype.suite.params.ModuleParametersHelper;
@@ -223,12 +227,14 @@
     private Map<String, LinkedHashSet<SuiteTestFilter>> mExcludeFiltersParsed = new LinkedHashMap<>();
     private List<File> mConfigPaths = new ArrayList<>();
     private Set<IAbi> mAbis = new LinkedHashSet<>();
+    private Set<DeviceFoldableState> mFoldableStates = new LinkedHashSet<>();
 
     /** {@inheritDoc} */
     @Override
     public LinkedHashMap<String, IConfiguration> loadTests() {
         try {
             File testsDir = getTestsDir();
+            mFoldableStates = getFoldableStates(getDevice());
             setupFilters(testsDir);
             mAbis = getAbis(getDevice());
 
@@ -324,6 +330,7 @@
             mModuleRepo.setOptionalParameterizedModules(mEnableOptionalParameter);
             mModuleRepo.setModuleParameter(mForceParameter);
             mModuleRepo.setExcludedModuleParameters(mExcludedModuleParameters);
+            mModuleRepo.setFoldableStates(mFoldableStates);
 
             List<File> testsDirectories = new ArrayList<>();
 
@@ -551,6 +558,21 @@
                         if (moduleParam.getValue() instanceof NegativeHandler) {
                             continue;
                         }
+                        if (moduleParam.getValue() instanceof FoldableExpandingHandler) {
+                            List<IModuleParameterHandler> foldableHandlers =
+                                    ((FoldableExpandingHandler) moduleParam.getValue())
+                                        .expandHandler(mFoldableStates);
+                            for (IModuleParameterHandler foldableHandler : foldableHandlers) {
+                                String paramModuleName =
+                                        String.format(
+                                                "%s[%s]", moduleName,
+                                                foldableHandler.getParameterIdentifier());
+                                mIncludeFilters.add(
+                                        new SuiteTestFilter(getRequestedAbi(), paramModuleName,
+                                                            mTestName).toString());
+                            }
+                            continue;
+                        }
                         String paramModuleName =
                                 String.format(
                                         "%s[%s]", moduleName,
@@ -660,4 +682,16 @@
         }
         return true;
     }
+
+    protected Set<DeviceFoldableState> getFoldableStates(ITestDevice device)
+            throws DeviceNotAvailableException {
+        if (device.getIDevice() instanceof StubDevice) {
+            return mFoldableStates;
+        }
+        if (!mFoldableStates.isEmpty()) {
+            return mFoldableStates;
+        }
+        mFoldableStates = device.getFoldableStates();
+        return mFoldableStates;
+    }
 }
diff --git a/src/com/android/tradefed/testtype/suite/SuiteModuleLoader.java b/src/com/android/tradefed/testtype/suite/SuiteModuleLoader.java
index ef47c10..035c405 100644
--- a/src/com/android/tradefed/testtype/suite/SuiteModuleLoader.java
+++ b/src/com/android/tradefed/testtype/suite/SuiteModuleLoader.java
@@ -24,6 +24,7 @@
 import com.android.tradefed.config.IConfigurationFactory;
 import com.android.tradefed.config.IDeviceConfiguration;
 import com.android.tradefed.config.OptionDef;
+import com.android.tradefed.device.DeviceFoldableState;
 import com.android.tradefed.error.HarnessRuntimeException;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.log.LogUtil.CLog;
@@ -33,6 +34,7 @@
 import com.android.tradefed.testtype.IRemoteTest;
 import com.android.tradefed.testtype.ITestFileFilterReceiver;
 import com.android.tradefed.testtype.ITestFilterReceiver;
+import com.android.tradefed.testtype.suite.params.FoldableExpandingHandler;
 import com.android.tradefed.testtype.suite.params.IModuleParameterHandler;
 import com.android.tradefed.testtype.suite.params.MainlineModuleHandler;
 import com.android.tradefed.testtype.suite.params.ModuleParameters;
@@ -87,6 +89,7 @@
     private boolean mAllowOptionalParameterizedModules = false;
     private ModuleParameters mForcedModuleParameter = null;
     private Set<ModuleParameters> mExcludedModuleParameters = new HashSet<>();
+    private Set<DeviceFoldableState> mFoldableStates = new LinkedHashSet<>();
     // Check the mainline parameter configured in a test config must end with .apk, .apks, or .apex.
     private static final Set<String> MAINLINE_PARAMETERS_TO_VALIDATE =
             new HashSet<>(Arrays.asList(".apk", ".apks", ".apex"));
@@ -151,6 +154,11 @@
         mExcludedModuleParameters = excludedParams;
     }
 
+    /** Sets the set of {@link DeviceFoldableState} that should be run. */
+    public final void setFoldableStates(Set<DeviceFoldableState> foldableStates) {
+        mFoldableStates = foldableStates;
+    }
+
     /** Main loading of configurations, looking into the specified files */
     public LinkedHashMap<String, IConfiguration> loadConfigsFromSpecifiedPaths(
             List<File> listConfigFiles,
@@ -678,7 +686,8 @@
         for (ModuleParameters moduleParameters : mExcludedModuleParameters) {
             expandedExcludedModuleParameters.addAll(
                     ModuleParametersHelper.resolveParam(
-                            moduleParameters, mAllowOptionalParameterizedModules).keySet());
+                            moduleParameters,
+                            mAllowOptionalParameterizedModules).keySet());
         }
 
         for (String p : parameters) {
@@ -688,7 +697,8 @@
             }
             Map<ModuleParameters, IModuleParameterHandler> suiteParams =
                     ModuleParametersHelper.resolveParam(
-                            p.toUpperCase(), mAllowOptionalParameterizedModules);
+                            ModuleParameters.valueOf(p.toUpperCase()),
+                            mAllowOptionalParameterizedModules);
             for (Entry<ModuleParameters, IModuleParameterHandler>  suiteParamEntry : suiteParams.entrySet()) {
                 ModuleParameters suiteParam = suiteParamEntry.getKey();
                 String family = suiteParam.getFamily();
@@ -708,7 +718,14 @@
                     continue;
                 }
 
-                params.add(suiteParamEntry.getValue());
+                if (suiteParamEntry instanceof FoldableExpandingHandler) {
+                    List<IModuleParameterHandler> foldableHandlers =
+                            ((FoldableExpandingHandler) suiteParamEntry)
+                                .expandHandler(mFoldableStates);
+                    params.addAll(foldableHandlers);
+                } else {
+                    params.add(suiteParamEntry.getValue());
+                }
             }
         }
         return params;
diff --git a/src/com/android/tradefed/testtype/suite/params/FoldableExpandingHandler.java b/src/com/android/tradefed/testtype/suite/params/FoldableExpandingHandler.java
new file mode 100644
index 0000000..2f7be94
--- /dev/null
+++ b/src/com/android/tradefed/testtype/suite/params/FoldableExpandingHandler.java
@@ -0,0 +1,58 @@
+/*
+ * 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.tradefed.testtype.suite.params;
+
+import com.android.tradefed.config.IConfiguration;
+import com.android.tradefed.device.DeviceFoldableState;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+
+/**
+ * A {@link IModuleParameterHandler} expanding into more for each non-primary foldable
+ * configuration.
+ */
+public class FoldableExpandingHandler implements IModuleParameterHandler {
+
+    @Override
+    public String getParameterIdentifier() {
+        throw new UnsupportedOperationException("Should never be called");
+    }
+
+    @Override
+    public void applySetup(IConfiguration moduleConfiguration) {
+        throw new UnsupportedOperationException("Should never be called");
+    }
+
+    public List<IModuleParameterHandler> expandHandler(Set<DeviceFoldableState> states) {
+        List<DeviceFoldableState> foldableList = new ArrayList<>(states);
+        Collections.sort(foldableList);
+
+        List<IModuleParameterHandler> foldableStateToRun = new ArrayList<>();
+        // If there is no or only one foldable state, there is no need to parameterize.
+        if (foldableList.isEmpty() || foldableList.size() == 1) {
+            return foldableStateToRun;
+        }
+        // Consider the lowest identifier the 'primary state'
+        for (int i = 1; i < foldableList.size(); i++) {
+            foldableStateToRun.add(new FoldableHandler(
+                    foldableList.get(i).toString(), foldableList.get(i).getIdentifier()));
+        }
+        return foldableStateToRun;
+    }
+}
diff --git a/src/com/android/tradefed/testtype/suite/params/ModuleParameters.java b/src/com/android/tradefed/testtype/suite/params/ModuleParameters.java
index e96b4dd..bfb7fd1 100644
--- a/src/com/android/tradefed/testtype/suite/params/ModuleParameters.java
+++ b/src/com/android/tradefed/testtype/suite/params/ModuleParameters.java
@@ -30,7 +30,11 @@
     // Multi-user
     MULTIUSER("multiuser", "multiuser_family"),
     RUN_ON_WORK_PROFILE("run_on_work_profile", "run_on_work_profile_family"),
-    RUN_ON_SECONDARY_USER("run_on_secondary_user", "run_on_secondary_user_family");
+    RUN_ON_SECONDARY_USER("run_on_secondary_user", "run_on_secondary_user_family"),
+
+    // Foldable mode
+    ALL_FOLDABLE_STATES("all_foldable_states", "foldable_family"),
+    NO_FOLDABLE_STATES("no_foldable_states", "foldable_family");
 
     public static final String INSTANT_APP_FAMILY = "instant_app_family";
     public static final String MULTI_ABI_FAMILY = "multi_abi_family";
diff --git a/src/com/android/tradefed/testtype/suite/params/ModuleParametersHelper.java b/src/com/android/tradefed/testtype/suite/params/ModuleParametersHelper.java
index 37d98d1..ae27248 100644
--- a/src/com/android/tradefed/testtype/suite/params/ModuleParametersHelper.java
+++ b/src/com/android/tradefed/testtype/suite/params/ModuleParametersHelper.java
@@ -38,6 +38,9 @@
                 ModuleParameters.RUN_ON_WORK_PROFILE, new RunOnWorkProfileParameterHandler());
         sHandlerMap.put(
                 ModuleParameters.RUN_ON_SECONDARY_USER, new RunOnSecondaryUserParameterHandler());
+
+        sHandlerMap.put(ModuleParameters.NO_FOLDABLE_STATES, new NegativeHandler());
+        sHandlerMap.put(ModuleParameters.ALL_FOLDABLE_STATES, new FoldableExpandingHandler());
     }
 
     private static Map<ModuleParameters, Set<ModuleParameters>> sGroupMap = new HashMap<>();
@@ -69,16 +72,6 @@
     }
 
     /**
-     * Resolve a {@link ModuleParameters} from its {@link String} representation.
-     *
-     * @see #resolveParam(ModuleParameters, boolean)
-     */
-    public static Map<ModuleParameters, IModuleParameterHandler> resolveParam(
-            String param, boolean withOptional) {
-        return resolveParam(ModuleParameters.valueOf(param.toUpperCase()), withOptional);
-    }
-
-    /**
      * Get the all {@link ModuleParameters} which are sub-params of a given {@link
      * ModuleParameters}.
      *
diff --git a/src/com/android/tradefed/util/LogcatEventParser.java b/src/com/android/tradefed/util/LogcatEventParser.java
index 66cd88d..bb516ef 100644
--- a/src/com/android/tradefed/util/LogcatEventParser.java
+++ b/src/com/android/tradefed/util/LogcatEventParser.java
@@ -41,7 +41,7 @@
  * <p>This class interprets logcat messages and can inform the listener of events in both a blocking
  * and polling fashion.
  */
-public class LogcatEventParser implements Closeable {
+class GenericLogcatEventParser<LogcatEventType> implements Closeable {
 
     // define a custom logcat parser category for events of interest
     private static final String CUSTOM_CATEGORY = "parserEvent";
@@ -58,7 +58,7 @@
     private boolean mCancelled = false;
 
     /** Struct to hold a logcat event with the event type and triggering logcat message */
-    public static class LogcatEvent {
+    public class LogcatEvent {
         private LogcatEventType eventType;
         private String msg;
 
@@ -121,7 +121,7 @@
      *
      * @param device to read logcat from
      */
-    public LogcatEventParser(ITestDevice device) {
+    public GenericLogcatEventParser(ITestDevice device) {
         mDevice = device;
     }
 
@@ -209,3 +209,9 @@
         }
     }
 }
+
+public class LogcatEventParser extends GenericLogcatEventParser<LogcatEventType> {
+    public LogcatEventParser(ITestDevice device) {
+        super(device);
+    }
+}
diff --git a/src/com/android/tradefed/util/LogcatEventType.java b/src/com/android/tradefed/util/LogcatEventType.java
index b5a9b34..b72fd87 100644
--- a/src/com/android/tradefed/util/LogcatEventType.java
+++ b/src/com/android/tradefed/util/LogcatEventType.java
@@ -27,6 +27,7 @@
     INSTALL_PROGRESS_UPDATE,
     UPDATE_COMPLETE,
     UPDATE_COMPLETE_NOT_ACTIVE,
+    UPDATE_CLEANED_UP,
     // error found in logcat output
     ERROR,
     // error found in logcat output, but doesn't necessarily indicate OTA failure. Should retry.
diff --git a/test_framework/com/android/tradefed/testtype/rust/RustBinaryHostTest.java b/test_framework/com/android/tradefed/testtype/rust/RustBinaryHostTest.java
index 0b31c30..269a96c 100644
--- a/test_framework/com/android/tradefed/testtype/rust/RustBinaryHostTest.java
+++ b/test_framework/com/android/tradefed/testtype/rust/RustBinaryHostTest.java
@@ -219,6 +219,7 @@
         }
 
         CLog.d("Running test with filter '%s'", filter);
+        getRunUtil().setWorkingDir(file.getParentFile());
         CommandResult result =
                 getRunUtil().runTimedCmd(mTestTimeout, commandLine.toArray(new String[0]));
         if (!CommandStatus.SUCCESS.equals(result.getStatus())) {