Snap for 7714994 from f1752a19f74f624bee52e7d35e32389b70e426e2 to sc-qpr1-release

Change-Id: I2c0ae41d5422760dd30c2f1481938ed6d5f0caaa
diff --git a/src/com/android/wallpaper/model/SetWallpaperViewModel.java b/src/com/android/wallpaper/model/SetWallpaperViewModel.java
new file mode 100644
index 0000000..efd55d3
--- /dev/null
+++ b/src/com/android/wallpaper/model/SetWallpaperViewModel.java
@@ -0,0 +1,83 @@
+/*
+ * 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.wallpaper.model;
+
+import android.util.Log;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.lifecycle.LiveData;
+import androidx.lifecycle.MutableLiveData;
+import androidx.lifecycle.ViewModel;
+import androidx.lifecycle.ViewModelProvider;
+
+import com.android.wallpaper.module.WallpaperPersister.Destination;
+import com.android.wallpaper.module.WallpaperPersister.SetWallpaperCallback;
+
+/**
+ * {@link ViewModel} class that keeps track of the status of the "Set wallpaper" operation.
+ */
+public class SetWallpaperViewModel extends ViewModel {
+
+    private static final String TAG = "SetWallpaperViewModel";
+
+    public enum SetWallpaperStatus {
+        UNKNOWN, PENDING, SUCCESS, ERROR
+    }
+
+    /**
+     * @return a {@link SetWallpaperCallback} to set as listener for
+     * {@link WallpaperSetter#setCurrentWallpaper} that will update the
+     * {@link SetWallpaperViewModel} obtained from the given provider.
+     */
+    public static SetWallpaperCallback getCallback(@NonNull ViewModelProvider provider) {
+        SetWallpaperViewModel viewModel = provider.get(SetWallpaperViewModel.class);
+        return new SetWallpaperCallback() {
+            @Override
+            public void onSuccess(WallpaperInfo wallpaperInfo) {
+                Log.d(TAG, "SetWallpaperCallback success");
+                viewModel.mStatus.setValue(SetWallpaperStatus.SUCCESS);
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                Log.w(TAG, "SetWallpaperCallback error", throwable);
+                viewModel.mStatus.setValue(SetWallpaperStatus.ERROR);
+            }
+        };
+    }
+
+    private final MutableLiveData<SetWallpaperStatus> mStatus = new MutableLiveData<>();
+
+    @Destination
+    private int mDestination;
+
+    public SetWallpaperViewModel() {
+        mStatus.setValue(SetWallpaperStatus.UNKNOWN);
+    }
+
+    public LiveData<SetWallpaperStatus> getStatus() {
+        return mStatus;
+    }
+
+    public int getDestination() {
+        return mDestination;
+    }
+
+    public void setDestination(int destination) {
+        mDestination = destination;
+    }
+}
diff --git a/src/com/android/wallpaper/module/WallpaperSetter.java b/src/com/android/wallpaper/module/WallpaperSetter.java
index 7f2be7c..dbb6c76 100644
--- a/src/com/android/wallpaper/module/WallpaperSetter.java
+++ b/src/com/android/wallpaper/module/WallpaperSetter.java
@@ -154,6 +154,7 @@
                                 if (event == Event.ON_DESTROY) {
                                     if (mProgressDialog != null) {
                                         mProgressDialog.dismiss();
+                                        mProgressDialog = null;
                                     }
                                 }
                             }
diff --git a/src/com/android/wallpaper/picker/ImagePreviewFragment.java b/src/com/android/wallpaper/picker/ImagePreviewFragment.java
index c41d394..1d7fa35 100755
--- a/src/com/android/wallpaper/picker/ImagePreviewFragment.java
+++ b/src/com/android/wallpaper/picker/ImagePreviewFragment.java
@@ -60,11 +60,10 @@
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.asset.CurrentWallpaperAssetVN;
-import com.android.wallpaper.model.WallpaperInfo;
+import com.android.wallpaper.model.SetWallpaperViewModel;
 import com.android.wallpaper.module.BitmapCropper;
 import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.WallpaperPersister.Destination;
-import com.android.wallpaper.module.WallpaperPersister.SetWallpaperCallback;
 import com.android.wallpaper.util.FullScreenAnimation;
 import com.android.wallpaper.util.ResourceUtils;
 import com.android.wallpaper.util.ScreenSizeCalculator;
@@ -480,17 +479,7 @@
     protected void setCurrentWallpaper(@Destination int destination) {
         mWallpaperSetter.setCurrentWallpaper(getActivity(), mWallpaper, mWallpaperAsset,
                 destination, mFullResImageView.getScale(), calculateCropRect(getContext()),
-                new SetWallpaperCallback() {
-                    @Override
-                    public void onSuccess(WallpaperInfo wallpaperInfo) {
-                        finishActivity(/* success= */ true);
-                    }
-
-                    @Override
-                    public void onError(@Nullable Throwable throwable) {
-                        showSetWallpaperErrorDialog(destination);
-                    }
-                });
+                SetWallpaperViewModel.getCallback(mViewModelProvider));
     }
 
     private void renderWorkspaceSurface() {
diff --git a/src/com/android/wallpaper/picker/LivePreviewFragment.java b/src/com/android/wallpaper/picker/LivePreviewFragment.java
index ec17767..b342e24 100644
--- a/src/com/android/wallpaper/picker/LivePreviewFragment.java
+++ b/src/com/android/wallpaper/picker/LivePreviewFragment.java
@@ -65,7 +65,7 @@
 
 import com.android.wallpaper.R;
 import com.android.wallpaper.compat.BuildCompat;
-import com.android.wallpaper.module.WallpaperPersister.SetWallpaperCallback;
+import com.android.wallpaper.model.SetWallpaperViewModel;
 import com.android.wallpaper.util.FullScreenAnimation;
 import com.android.wallpaper.util.ResourceUtils;
 import com.android.wallpaper.util.ScreenSizeCalculator;
@@ -444,17 +444,7 @@
     @Override
     protected void setCurrentWallpaper(int destination) {
         mWallpaperSetter.setCurrentWallpaper(getActivity(), mWallpaper, null,
-                destination, 0, null, new SetWallpaperCallback() {
-                    @Override
-                    public void onSuccess(com.android.wallpaper.model.WallpaperInfo wallpaperInfo) {
-                        finishActivity(/* success= */ true);
-                    }
-
-                    @Override
-                    public void onError(@Nullable Throwable throwable) {
-                        showSetWallpaperErrorDialog(destination);
-                    }
-                });
+                destination, 0, null, SetWallpaperViewModel.getCallback(mViewModelProvider));
     }
 
     @Nullable
diff --git a/src/com/android/wallpaper/picker/PreviewFragment.java b/src/com/android/wallpaper/picker/PreviewFragment.java
index 74bdd6f..d0626f0 100755
--- a/src/com/android/wallpaper/picker/PreviewFragment.java
+++ b/src/com/android/wallpaper/picker/PreviewFragment.java
@@ -23,6 +23,7 @@
 import android.content.Intent;
 import android.content.res.Resources.NotFoundException;
 import android.os.Bundle;
+import android.os.Handler;
 import android.util.Log;
 import android.view.LayoutInflater;
 import android.view.SurfaceView;
@@ -39,9 +40,11 @@
 import androidx.annotation.LayoutRes;
 import androidx.annotation.Nullable;
 import androidx.fragment.app.FragmentActivity;
+import androidx.lifecycle.ViewModelProvider;
 
 import com.android.wallpaper.R;
 import com.android.wallpaper.model.LiveWallpaperInfo;
+import com.android.wallpaper.model.SetWallpaperViewModel;
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.Injector;
 import com.android.wallpaper.module.InjectorProvider;
@@ -80,7 +83,6 @@
      * wallpaper with pan and crop position to the device.
      */
     static final int MODE_CROP_AND_SET_WALLPAPER = 1;
-    private Optional<Integer> mLastSelectedTabPositionOptional = Optional.empty();
 
     /**
      * Possible preview modes for the fragment.
@@ -139,6 +141,9 @@
     @PreviewMode protected int mPreviewMode;
     protected boolean mViewAsHome;
 
+    protected SetWallpaperViewModel mSetWallpaperViewModel;
+    protected ViewModelProvider mViewModelProvider;
+    private Optional<Integer> mLastSelectedTabPositionOptional = Optional.empty();
     private OnBackPressedCallback mOnBackPressedCallback;
 
     /**
@@ -166,6 +171,9 @@
         mWallpaperSetter = new WallpaperSetter(injector.getWallpaperPersister(appContext),
                 injector.getPreferences(appContext), mUserEventLogger, mTestingModeEnabled);
 
+        mViewModelProvider = new ViewModelProvider(requireActivity());
+        mSetWallpaperViewModel = mViewModelProvider.get(SetWallpaperViewModel.class);
+
         Activity activity = getActivity();
         List<String> attributions = getAttributions(activity);
         if (attributions.size() > 0 && attributions.get(0) != null) {
@@ -284,12 +292,28 @@
                     requireFragmentManager(), TAG_SET_WALLPAPER_ERROR_DIALOG_FRAGMENT);
             mStagedSetWallpaperErrorDialogFragment = null;
         }
+
+        mSetWallpaperViewModel.getStatus().observe(requireActivity(), setWallpaperStatus -> {
+            switch (setWallpaperStatus) {
+                case SUCCESS:
+                    // Give a few millis before finishing to allow for the dialog dismiss
+                    // and animations to finish
+                    Handler.getMain().postDelayed(() -> finishActivity(true), 300);
+                    break;
+                case ERROR:
+                    showSetWallpaperErrorDialog(mSetWallpaperViewModel.getDestination());
+                    break;
+                default:
+                    // Do nothing in this case, either status is pending, or unknown
+            }
+        });
     }
 
     protected abstract boolean isLoaded();
 
     @Override
     public void onSet(int destination) {
+        mSetWallpaperViewModel.setDestination(destination);
         setCurrentWallpaper(destination);
     }
 
@@ -300,6 +324,7 @@
 
     @Override
     public void onClickTryAgain(@Destination int wallpaperDestination) {
+        mSetWallpaperViewModel.setDestination(wallpaperDestination);
         setCurrentWallpaper(wallpaperDestination);
     }
 
@@ -375,8 +400,8 @@
             }
             activity.setResult(Activity.RESULT_OK);
         }
-        activity.overridePendingTransition(R.anim.fade_in, R.anim.fade_out);
         activity.finish();
+        activity.overridePendingTransition(R.anim.fade_in, R.anim.fade_out);
     }
 
     protected void showSetWallpaperErrorDialog(@Destination int wallpaperDestination) {