Add timeout for get account name in photo picker

Test: BannerControllerTest#testCloudProviderSlowQueryFallback
Bug: 294365231
Change-Id: I495ca99484c6bec6e58256da6dae3ede1cb385cf
diff --git a/src/com/android/providers/media/photopicker/ui/settings/SettingsCloudMediaViewModel.java b/src/com/android/providers/media/photopicker/ui/settings/SettingsCloudMediaViewModel.java
index 346eed3..ca9be63 100644
--- a/src/com/android/providers/media/photopicker/ui/settings/SettingsCloudMediaViewModel.java
+++ b/src/com/android/providers/media/photopicker/ui/settings/SettingsCloudMediaViewModel.java
@@ -56,6 +56,7 @@
 public class SettingsCloudMediaViewModel extends ViewModel {
     static final String NONE_PREF_KEY = "none";
     private static final String TAG = "SettingsFragVM";
+    private static final long GET_ACCOUNT_NAME_TIMEOUT_IN_MILLIS = 10000L;
 
     @NonNull
     private final Context mContext;
@@ -196,7 +197,8 @@
         } else {
             try {
                 final String accountName = getCloudMediaAccountName(
-                        mUserId.getContentResolver(mContext), currentProviderAuthority);
+                        mUserId.getContentResolver(mContext), currentProviderAuthority,
+                        GET_ACCOUNT_NAME_TIMEOUT_IN_MILLIS);
                 return new CloudMediaProviderAccount(currentProviderAuthority, accountName);
             } catch (Exception e) {
                 Log.w(TAG, "Failed to fetch account name from the cloud media provider.", e);
diff --git a/src/com/android/providers/media/photopicker/util/CloudProviderUtils.java b/src/com/android/providers/media/photopicker/util/CloudProviderUtils.java
index 741755b..dbddceb 100644
--- a/src/com/android/providers/media/photopicker/util/CloudProviderUtils.java
+++ b/src/com/android/providers/media/photopicker/util/CloudProviderUtils.java
@@ -28,8 +28,12 @@
 import static android.provider.MediaStore.PICKER_MEDIA_INIT_CALL;
 import static android.provider.MediaStore.SET_CLOUD_PROVIDER_CALL;
 
-import static java.util.Collections.emptyList;
+import static com.android.providers.media.PickerUriResolver.getMediaCollectionInfoUri;
 
+import static java.util.Collections.emptyList;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+
+import android.annotation.DurationMillisLong;
 import android.content.ContentProviderClient;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -55,6 +59,9 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Utility methods for retrieving available and/or allowlisted Cloud Providers.
@@ -248,23 +255,26 @@
     }
 
     /**
+     * @param resolver                    {@link ContentResolver} for the related user
+     * @param cloudMediaProviderAuthority authority {@link String} of the {@link CloudMediaProvider}
+     * @param timeout                     timeout in milliseconds for this query (<= 0 for timeout)
      * @return the current cloud media account name for the {@link CloudMediaProvider} with the
      *         given {@code cloudMediaProviderAuthority}.
      */
     @Nullable
     public static String getCloudMediaAccountName(@NonNull ContentResolver resolver,
-            @Nullable String cloudMediaProviderAuthority) {
+            @Nullable String cloudMediaProviderAuthority, @DurationMillisLong long timeout)
+            throws ExecutionException, InterruptedException, TimeoutException {
         if (cloudMediaProviderAuthority == null) {
             return null;
         }
 
-        try (ContentProviderClient client =
-                     resolver.acquireContentProviderClient(cloudMediaProviderAuthority)) {
-            final Bundle out = client.call(METHOD_GET_MEDIA_COLLECTION_INFO, /* arg */ null,
-                    /* extras */ null);
-            return out.getString(ACCOUNT_NAME);
-        } catch (RemoteException e) {
-            throw e.rethrowAsRuntimeException();
-        }
+        CompletableFuture<String> future = CompletableFuture.supplyAsync(() -> {
+            final Bundle out = resolver.call(getMediaCollectionInfoUri(cloudMediaProviderAuthority),
+                    METHOD_GET_MEDIA_COLLECTION_INFO, /* arg */ null, /* extras */ null);
+            return (out == null) ? null : out.getString(ACCOUNT_NAME);
+        });
+
+        return (timeout > 0) ? future.get(timeout, MILLISECONDS) : future.get();
     }
 }
diff --git a/src/com/android/providers/media/photopicker/viewmodel/BannerController.java b/src/com/android/providers/media/photopicker/viewmodel/BannerController.java
index 5b070a3..b2fc873 100644
--- a/src/com/android/providers/media/photopicker/viewmodel/BannerController.java
+++ b/src/com/android/providers/media/photopicker/viewmodel/BannerController.java
@@ -44,6 +44,8 @@
 import java.io.FileOutputStream;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Banner Controller to store and handle the banner data per user for
@@ -64,6 +66,7 @@
      * {@link android.provider.CloudMediaProvider}.
      */
     private static final String ACCOUNT_NAME = "account_name";
+    private static final long GET_ACCOUNT_NAME_TIMEOUT_IN_MILLIS = 100L;
 
     private final Context mContext;
     private final UserHandle mUserHandle;
@@ -130,7 +133,8 @@
      * block the UI thread on the heavy Binder calls to fetch the cloud media provider info.
      */
     private void initialise() {
-        final String cmpAuthority, cmpAccountName;
+        String cmpAuthority = null, cmpAccountName = null;
+        mCmpLabel = null;
         // TODO(b/245746037): Remove try-catch for the RuntimeException.
         //  Under the hood MediaStore.getCurrentCloudProvider() makes an IPC call to the primary
         //  MediaProvider process, where we currently perform a UID check (making sure that
@@ -149,14 +153,17 @@
                     UserId.of(mUserHandle).getContentResolver(mContext);
             cmpAuthority = getCurrentCloudProvider(contentResolver);
             mCmpLabel = getProviderLabelForUser(mContext, mUserHandle, cmpAuthority);
-            cmpAccountName = getCloudMediaAccountName(contentResolver, cmpAuthority);
+            cmpAccountName = getCloudMediaAccountName(contentResolver, cmpAuthority,
+                    GET_ACCOUNT_NAME_TIMEOUT_IN_MILLIS);
 
             // Not logging the account name due to privacy concerns
             Log.d(TAG, "Current CloudMediaProvider authority: " + cmpAuthority + ", label: "
                     + mCmpLabel);
-        } catch (PackageManager.NameNotFoundException | RuntimeException e) {
+        } catch (PackageManager.NameNotFoundException | RuntimeException | ExecutionException
+                | InterruptedException | TimeoutException e) {
             Log.w(TAG, "Could not fetch the current CloudMediaProvider", e);
-            resetToDefault();
+            updateCloudProviderDataMap(cmpAuthority, cmpAccountName);
+            clearBanners();
             return;
         }
 
@@ -213,15 +220,6 @@
     }
 
     /**
-     * Reset all the controller data to their default values.
-     */
-    private void resetToDefault() {
-        mCloudProviderDataMap.clear();
-        mCmpLabel = null;
-        clearBanners();
-    }
-
-    /**
      * Clear all banners
      *
      * Reset all should show banner {@code boolean} values to {@code false}.
@@ -385,6 +383,12 @@
 
     private void persistCloudProviderInfo(@Nullable String cmpAuthority,
             @Nullable String cmpAccountName) {
+        updateCloudProviderDataMap(cmpAuthority, cmpAccountName);
+        updateCloudProviderDataFile();
+    }
+
+    private void updateCloudProviderDataMap(@Nullable String cmpAuthority,
+            @Nullable String cmpAccountName) {
         mCloudProviderDataMap.clear();
         if (cmpAuthority != null) {
             mCloudProviderDataMap.put(AUTHORITY, cmpAuthority);
@@ -392,8 +396,6 @@
         if (cmpAccountName != null) {
             mCloudProviderDataMap.put(ACCOUNT_NAME, cmpAccountName);
         }
-
-        updateCloudProviderDataFile();
     }
 
     @VisibleForTesting
diff --git a/tests/src/com/android/providers/media/IsolatedContext.java b/tests/src/com/android/providers/media/IsolatedContext.java
index 8a87eb5..dc52fcd 100644
--- a/tests/src/com/android/providers/media/IsolatedContext.java
+++ b/tests/src/com/android/providers/media/IsolatedContext.java
@@ -31,8 +31,10 @@
 import android.test.mock.MockContentResolver;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.providers.media.cloudproviders.CloudProviderPrimary;
+import com.android.providers.media.cloudproviders.FlakyCloudProvider;
 import com.android.providers.media.photopicker.PhotoPickerProvider;
 import com.android.providers.media.photopicker.PickerSyncController;
 import com.android.providers.media.util.FileUtils;
@@ -48,6 +50,7 @@
     private final MockContentResolver mResolver;
     private final MediaProvider mMediaProvider;
     private final UserHandle mUserHandle;
+    private final FlakyCloudProvider mFlakyCloudProvider;
 
     public IsolatedContext(Context base, String tag, boolean asFuseThread) {
         this(base, tag, asFuseThread, base.getUser());
@@ -88,6 +91,9 @@
         final CloudMediaProvider cmp = new CloudProviderPrimary();
         attachInfoAndAddProvider(base, cmp, CloudProviderPrimary.AUTHORITY);
 
+        mFlakyCloudProvider = new FlakyCloudProvider();
+        attachInfoAndAddProvider(base, mFlakyCloudProvider, FlakyCloudProvider.AUTHORITY);
+
         MediaStore.waitForIdle(mResolver);
     }
 
@@ -162,4 +168,13 @@
         }
     }
 
+    @VisibleForTesting
+    public void setFlakyCloudProviderToFlakeInTheNextRequest() {
+        mFlakyCloudProvider.setToFlakeInTheNextRequest();
+    }
+
+    @VisibleForTesting
+    public void resetFlakyCloudProviderToNotFlakeInTheNextRequest() {
+        mFlakyCloudProvider.resetToNotFlakeInTheNextRequest();
+    }
 }
diff --git a/tests/src/com/android/providers/media/cloudproviders/FlakyCloudProvider.java b/tests/src/com/android/providers/media/cloudproviders/FlakyCloudProvider.java
index d4c3dec..2d20574 100644
--- a/tests/src/com/android/providers/media/cloudproviders/FlakyCloudProvider.java
+++ b/tests/src/com/android/providers/media/cloudproviders/FlakyCloudProvider.java
@@ -20,6 +20,8 @@
 
 import static com.android.providers.media.PickerProviderMediaGenerator.MediaGenerator;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+
 import android.content.res.AssetFileDescriptor;
 import android.database.Cursor;
 import android.graphics.Point;
@@ -27,6 +29,9 @@
 import android.os.CancellationSignal;
 import android.os.ParcelFileDescriptor;
 import android.provider.CloudMediaProvider;
+import android.util.Log;
+
+import androidx.annotation.VisibleForTesting;
 
 import com.android.providers.media.PickerProviderMediaGenerator;
 import com.android.providers.media.photopicker.data.CloudProviderQueryExtras;
@@ -41,24 +46,28 @@
  * out of every three requests.
  */
 public class FlakyCloudProvider extends CloudMediaProvider {
+    private static final String TAG = "FlakyCloudProvider";
     public static final String AUTHORITY =
             "com.android.providers.media.photopicker.tests.cloud_flaky";
+    public static final String ACCOUNT_NAME = "test_account@flakyCloudProvider";
+    private static final int INITIAL_REQUEST_COUNT = 0;
+    private static final int REQUEST_COUNT_FOR_NEXT_ONE_TO_FLAKE = 2;
 
     private final MediaGenerator mMediaGenerator =
             PickerProviderMediaGenerator.getMediaGenerator(AUTHORITY);
 
-    private int mRequestCount = 0;
+    private int mRequestCount = INITIAL_REQUEST_COUNT;
 
     /** Determines if the current request should flake. */
     private boolean shouldFlake() {
 
         // Always succeed on the first request.
-        if (++mRequestCount < 2) {
+        if (++mRequestCount < REQUEST_COUNT_FOR_NEXT_ONE_TO_FLAKE) {
             return false;
         }
 
-        if (mRequestCount >= 3) {
-            mRequestCount = 0;
+        if (mRequestCount > REQUEST_COUNT_FOR_NEXT_ONE_TO_FLAKE) {
+            mRequestCount = INITIAL_REQUEST_COUNT;
         }
 
         return true;
@@ -66,6 +75,7 @@
 
     @Override
     public boolean onCreate() {
+        mMediaGenerator.setAccountInfo(ACCOUNT_NAME, /* configIntent= */ null);
         return true;
     }
 
@@ -137,6 +147,24 @@
 
     @Override
     public Bundle onGetMediaCollectionInfo(Bundle extras) {
+        if (shouldFlake()) {
+            try {
+                MILLISECONDS.sleep(/* timeout= */ 200L);
+            } catch (InterruptedException e) {
+                Log.d(TAG, "Error while sleep when should flake on get media collection info.", e);
+            }
+        }
+
         return mMediaGenerator.getMediaCollectionInfo();
     }
+
+    @VisibleForTesting
+    public void resetToNotFlakeInTheNextRequest() {
+        mRequestCount = INITIAL_REQUEST_COUNT;
+    }
+
+    @VisibleForTesting
+    public void setToFlakeInTheNextRequest() {
+        mRequestCount = REQUEST_COUNT_FOR_NEXT_ONE_TO_FLAKE;
+    }
 }
diff --git a/tests/src/com/android/providers/media/photopicker/viewmodel/BannerControllerTest.java b/tests/src/com/android/providers/media/photopicker/viewmodel/BannerControllerTest.java
index d1f8a58..05ccbbe 100644
--- a/tests/src/com/android/providers/media/photopicker/viewmodel/BannerControllerTest.java
+++ b/tests/src/com/android/providers/media/photopicker/viewmodel/BannerControllerTest.java
@@ -16,19 +16,30 @@
 
 package com.android.providers.media.photopicker.viewmodel;
 
+import static android.provider.MediaStore.AUTHORITY;
+import static android.provider.MediaStore.getCurrentCloudProvider;
+
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.providers.media.photopicker.util.CloudProviderUtils.persistSelectedProvider;
+
 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 android.content.ContentProviderClient;
+import android.content.ContentResolver;
 import android.content.Context;
+import android.os.RemoteException;
 
+import androidx.annotation.Nullable;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.providers.media.IsolatedContext;
 import com.android.providers.media.TestConfigStore;
+import com.android.providers.media.cloudproviders.FlakyCloudProvider;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -36,37 +47,37 @@
 
 @RunWith(AndroidJUnit4.class)
 public class BannerControllerTest {
-    private BannerController mBannerController;
+    private static final Context sTargetContext = getInstrumentation().getTargetContext();
+    private static final String TEST_PACKAGE_NAME = "com.android.providers.media.tests";
     private static final String CMP_AUTHORITY = "authority";
     private static final String CMP_ACCOUNT_NAME = "account_name";
-    private static final String TAG = "BannerControllerTest";
+
+    private IsolatedContext mIsolatedContext;
+    private ContentResolver mContentResolver;
+    private BannerController mBannerController;
 
     @Before
-    public void setUp() {
-        final Context context = new IsolatedContext(
-                getInstrumentation().getTargetContext(), TAG, /* asFuseThread= */ false);
+    public void setUp() throws RemoteException {
         final TestConfigStore configStore = new TestConfigStore();
+        configStore.enableCloudMediaFeatureAndSetAllowedCloudProviderPackages(TEST_PACKAGE_NAME);
 
-        mBannerController = new BannerController(context, context.getUser(), configStore) {
-            @Override
-            void updateCloudProviderDataFile() {
-                // No-op
-            }
+        mIsolatedContext = new IsolatedContext(sTargetContext, /* tag= */ "databases",
+                /* asFuseThread= */ false, sTargetContext.getUser(), configStore);
+        mContentResolver = mIsolatedContext.getContentResolver();
 
-            @Override
-            boolean areCloudProviderOptionsAvailable() {
-                return true;
-            }
-        };
+        setCloudProvider(/* authority= */ null);
+
+        mBannerController =
+                new BannerController(mIsolatedContext, mIsolatedContext.getUser(), configStore) {
+                    @Override
+                    void updateCloudProviderDataFile() {
+                        // No-op
+                    }
+                };
 
         assertNull(mBannerController.getCloudMediaProviderAuthority());
         assertNull(mBannerController.getCloudMediaProviderLabel());
         assertNull(mBannerController.getCloudMediaProviderAccountName());
-
-        assertFalse(mBannerController.shouldShowCloudMediaAvailableBanner());
-        assertFalse(mBannerController.shouldShowAccountUpdatedBanner());
-        assertFalse(mBannerController.shouldShowChooseAccountBanner());
-        assertFalse(mBannerController.shouldShowChooseAppBanner());
     }
 
     @Test
@@ -150,4 +161,51 @@
         assertFalse(mBannerController.shouldShowChooseAccountBanner());
         assertFalse(mBannerController.shouldShowChooseAppBanner());
     }
+
+    @Test
+    public void testCloudProviderSlowQueryFallback() throws RemoteException {
+        setCloudProvider(FlakyCloudProvider.AUTHORITY);
+
+        // Test for fast query
+        mIsolatedContext.resetFlakyCloudProviderToNotFlakeInTheNextRequest();
+        mBannerController.onChangeCloudMediaInfo(
+                /* cmpAuthority */ null, /* cmpAccountName */ null);
+        mBannerController.reset();
+
+        assertEquals(FlakyCloudProvider.AUTHORITY,
+                mBannerController.getCloudMediaProviderAuthority());
+        assertEquals(FlakyCloudProvider.ACCOUNT_NAME,
+                mBannerController.getCloudMediaProviderAccountName());
+
+        assertTrue(mBannerController.shouldShowCloudMediaAvailableBanner());
+        assertFalse(mBannerController.shouldShowAccountUpdatedBanner());
+        assertFalse(mBannerController.shouldShowChooseAccountBanner());
+        assertFalse(mBannerController.shouldShowChooseAppBanner());
+
+        // Test for slow query
+        mIsolatedContext.setFlakyCloudProviderToFlakeInTheNextRequest();
+        mBannerController.onChangeCloudMediaInfo(
+                /* cmpAuthority */ null, /* cmpAccountName */ null);
+        mBannerController.reset();
+
+        assertEquals(FlakyCloudProvider.AUTHORITY,
+                mBannerController.getCloudMediaProviderAuthority());
+        assertNull(mBannerController.getCloudMediaProviderAccountName());
+
+        assertFalse(mBannerController.shouldShowCloudMediaAvailableBanner());
+        assertFalse(mBannerController.shouldShowAccountUpdatedBanner());
+        assertFalse(mBannerController.shouldShowChooseAccountBanner());
+        assertFalse(mBannerController.shouldShowChooseAppBanner());
+    }
+
+    private void setCloudProvider(@Nullable String authority) throws RemoteException {
+        final ContentProviderClient client =
+                mContentResolver.acquireContentProviderClient(AUTHORITY);
+        assertNotNull(client);
+
+        persistSelectedProvider(client, authority);
+
+        final String actualAuthority = getCurrentCloudProvider(mContentResolver);
+        assertEquals(authority, actualAuthority);
+    }
 }