DLC: Allow pseudo lock/unlock in the unprovisioned state

Bug: 316024740
Test: [main] atest CtsDeviceLockTestCases
Test: [android14-tests-dev] cts-tradefed run cts -m CtsDeviceLockTestCases

Change-Id: I43ff468fc14a7791a4ee6012ab4a315061c2a2c0
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/policy/DeviceStateControllerImpl.java b/DeviceLockController/src/com/android/devicelockcontroller/policy/DeviceStateControllerImpl.java
index 0c4b303..b860fe1 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/policy/DeviceStateControllerImpl.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/policy/DeviceStateControllerImpl.java
@@ -23,6 +23,7 @@
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionEvent.PROVISION_SUCCESS;
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.KIOSK_PROVISIONED;
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.PROVISION_SUCCEEDED;
+import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.UNPROVISIONED;
 
 import com.android.devicelockcontroller.storage.GlobalParametersClient;
 
@@ -38,6 +39,10 @@
     private final DevicePolicyController mPolicyController;
     private final GlobalParametersClient mGlobalParametersClient;
     private final Executor mExecutor;
+    // Used to exercising APIs under CTS without actually applying any policies.
+    // This is not persistent across controller restarts, but should be good enough for the
+    // intended purpose.
+    private volatile @DeviceState int mPseudoDeviceState;
 
     public DeviceStateControllerImpl(DevicePolicyController policyController,
             ProvisionStateController provisionStateController, Executor executor) {
@@ -45,6 +50,7 @@
         mProvisionStateController = provisionStateController;
         mGlobalParametersClient = GlobalParametersClient.getInstance();
         mExecutor = executor;
+        mPseudoDeviceState = UNDEFINED;
     }
 
     @Override
@@ -79,6 +85,13 @@
                                 mProvisionStateController.setNextStateForEvent(PROVISION_SUCCESS);
                     } else if (provisionState == PROVISION_SUCCEEDED) {
                         maybeSetProvisioningSuccess = Futures.immediateVoidFuture();
+                    } else if (provisionState == UNPROVISIONED && (deviceState == LOCKED
+                        || deviceState == UNLOCKED)) {
+                        // During normal operation, we should not get lock/unlock requests in
+                        // the UNPROVISIONED state. Used for CTS compliance.
+                        mPseudoDeviceState = deviceState;
+                        // Do not apply any policies
+                        return Futures.immediateVoidFuture();
                     } else {
                         throw new RuntimeException("User has not been provisioned!");
                     }
@@ -99,14 +112,23 @@
 
     @Override
     public ListenableFuture<Boolean> isLocked() {
-        return Futures.transform(mGlobalParametersClient.getDeviceState(),
-                s -> {
-                    if (s == UNDEFINED) {
-                        throw new IllegalStateException("isLocked called before setting the "
-                                + "locked state (lockDevice/unlockDevice)");
+        return Futures.transformAsync(mProvisionStateController.getState(),
+                provisionState -> {
+                    if (provisionState == UNDEFINED) {
+                        // Used for CTS compliance.
+                        return Futures.immediateFuture(mPseudoDeviceState == LOCKED);
+                    } else {
+                        return Futures.transform(mGlobalParametersClient.getDeviceState(),
+                                s -> {
+                                    if (s == UNDEFINED) {
+                                        throw new IllegalStateException("isLocked called before "
+                                                + "setting the locked state "
+                                                + "(lockDevice/unlockDevice)");
+                                    }
+                                    return s == LOCKED;
+                                }, mExecutor);
                     }
-                    return s == LOCKED;
-            }, MoreExecutors.directExecutor());
+                }, mExecutor);
     }
 
     private ListenableFuture<Boolean> isCleared() {
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/DeviceStateControllerImplTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/DeviceStateControllerImplTest.java
index ba7dd0f..4d8fe73 100644
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/DeviceStateControllerImplTest.java
+++ b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/DeviceStateControllerImplTest.java
@@ -63,16 +63,15 @@
     }
 
     @Test
-    public void lockDevice_withUnprovisionedState_shouldThrowException()
+    public void lockDevice_withUnprovisionedState_shouldPseudoLockDevice()
             throws ExecutionException, InterruptedException {
         when(mMockProvisionStateController.getState()).thenReturn(
                 Futures.immediateFuture(ProvisionState.UNPROVISIONED));
+        mDeviceStateController.lockDevice().get();
 
-        ExecutionException thrown = assertThrows(ExecutionException.class,
-                () -> mDeviceStateController.lockDevice().get());
-        assertThat(thrown).hasCauseThat().isInstanceOf(RuntimeException.class);
-        assertThat(thrown).hasMessageThat().contains(USER_HAS_NOT_BEEN_PROVISIONED);
+        assertThat(mDeviceStateController.isLocked().get()).isTrue();
 
+        // Should not have changed the real device state
         assertThat(GlobalParametersClient.getInstance().getDeviceState().get()).isEqualTo(
                 DeviceState.UNDEFINED);
     }
@@ -183,16 +182,15 @@
     }
 
     @Test
-    public void unlockDevice_withUnprovisionedState_shouldThrowException()
+    public void unlockDevice_withUnprovisionedState_shouldPseudoUnlock()
             throws ExecutionException, InterruptedException {
         when(mMockProvisionStateController.getState()).thenReturn(
                 Futures.immediateFuture(ProvisionState.UNPROVISIONED));
+        mDeviceStateController.unlockDevice().get();
 
-        ExecutionException thrown = assertThrows(ExecutionException.class,
-                () -> mDeviceStateController.unlockDevice().get());
-        assertThat(thrown).hasCauseThat().isInstanceOf(RuntimeException.class);
-        assertThat(thrown).hasMessageThat().contains(USER_HAS_NOT_BEEN_PROVISIONED);
+        assertThat(mDeviceStateController.isLocked().get()).isFalse();
 
+        // Should not have changed the real device state
         assertThat(GlobalParametersClient.getInstance().getDeviceState().get()).isEqualTo(
                 DeviceState.UNDEFINED);
     }
diff --git a/tests/cts/src/com/android/cts/devicelock/DeviceLockManagerTest.java b/tests/cts/src/com/android/cts/devicelock/DeviceLockManagerTest.java
index ae0d740..a284059 100644
--- a/tests/cts/src/com/android/cts/devicelock/DeviceLockManagerTest.java
+++ b/tests/cts/src/com/android/cts/devicelock/DeviceLockManagerTest.java
@@ -254,15 +254,15 @@
         try {
             addFinancedDeviceKioskRole();
 
-            assertThrows(ExecutionException.class,
-                    () -> getLockDeviceFuture().get(TIMEOUT, TimeUnit.SECONDS));
+            getLockDeviceFuture().get(TIMEOUT, TimeUnit.SECONDS);
 
-            assertThrows(ExecutionException.class,
-                    () -> getUnlockDeviceFuture().get(TIMEOUT, TimeUnit.SECONDS));
+            boolean locked = getIsDeviceLockedFuture().get(TIMEOUT, TimeUnit.SECONDS);
+            assertThat(locked).isTrue();
 
-            ExecutionException executionException = assertThrows(ExecutionException.class,
-                    () -> getIsDeviceLockedFuture().get(TIMEOUT, TimeUnit.SECONDS));
-            assertThat(executionException).hasCauseThat().isInstanceOf(IllegalStateException.class);
+            getUnlockDeviceFuture().get(TIMEOUT, TimeUnit.SECONDS);
+
+            locked = getIsDeviceLockedFuture().get(TIMEOUT, TimeUnit.SECONDS);
+            assertThat(locked).isFalse();
         } finally {
             removeFinancedDeviceKioskRole();
         }