Add a unittest to cover the retry of load escrow data

The load escrow data is async. Add a unit test to cover the retry of
temporarily failure.

Bug: 172780686
Test: atest FrameworksServicesTests:RebootEscrowManagerTests
Change-Id: I63b264a5f9b68039868fcf962f4654c8d587985d

1 files changed, 54 insertions, 0 deletions

diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
index a896f1b0d60f..b51f4df43259 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -44,6 +44,7 @@ import android.content.ContextWrapper;
 import android.content.pm.UserInfo;
 import android.hardware.rebootescrow.IRebootEscrow;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
 import android.os.UserManager;
@@ -62,6 +63,7 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 
 import java.io.File;
+import java.io.IOException;
 import java.util.ArrayList;
 
 import javax.crypto.SecretKey;
@@ -181,6 +183,18 @@ public class RebootEscrowManagerTests {
         }
 
         @Override
+        public int getLoadEscrowDataRetryLimit() {
+            // Try two times
+            return 2;
+        }
+
+        @Override
+        public int getLoadEscrowDataRetryIntervalSeconds() {
+            // Retry in 1 seconds
+            return 1;
+        }
+
+        @Override
         public void reportMetric(boolean success) {
             mInjected.reportMetric(success);
         }
@@ -448,6 +462,46 @@ public class RebootEscrowManagerTests {
     }
 
     @Test
+    public void loadRebootEscrowDataIfAvailable_ServerBased_RetrySuccess() throws Exception {
+        setServerBasedRebootEscrowProvider();
+
+        when(mInjected.getBootCount()).thenReturn(0);
+        RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
+        mService.setRebootEscrowListener(mockListener);
+        mService.prepareRebootEscrow();
+
+        clearInvocations(mServiceConnection);
+        mService.callToRebootEscrowIfNeeded(PRIMARY_USER_ID, FAKE_SP_VERSION, FAKE_AUTH_TOKEN);
+        verify(mockListener).onPreparedForReboot(eq(true));
+        verify(mServiceConnection, never()).wrapBlob(any(), anyLong(), anyLong());
+
+        // Use x -> x for both wrap & unwrap functions.
+        when(mServiceConnection.wrapBlob(any(), anyLong(), anyLong()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+        assertTrue(mService.armRebootEscrowIfNeeded());
+        verify(mServiceConnection).wrapBlob(any(), anyLong(), anyLong());
+        assertTrue(mStorage.hasRebootEscrowServerBlob());
+
+        // pretend reboot happens here
+        when(mInjected.getBootCount()).thenReturn(1);
+        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
+        doNothing().when(mInjected).reportMetric(metricsSuccessCaptor.capture());
+
+        when(mServiceConnection.unwrap(any(), anyLong()))
+                .thenThrow(new IOException())
+                .thenAnswer(invocation -> invocation.getArgument(0));
+
+        HandlerThread thread = new HandlerThread("RebootEscrowManagerTest");
+        thread.start();
+        mService.loadRebootEscrowDataIfAvailable(new Handler(thread.getLooper()));
+        // Sleep 5s for the retry to complete
+        Thread.sleep(5 * 1000);
+        verify(mServiceConnection, times(2)).unwrap(any(), anyLong());
+        assertTrue(metricsSuccessCaptor.getValue());
+        verify(mKeyStoreManager).clearKeyStoreEncryptionKey();
+    }
+
+    @Test
     public void loadRebootEscrowDataIfAvailable_TooManyBootsInBetween_NoMetrics() throws Exception {
         when(mInjected.getBootCount()).thenReturn(0);