1 files changed, 26 insertions, 7 deletions
diff --git a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
index a1623d18..5d551a6a 100644
--- a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
+++ b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
@@ -20,11 +20,13 @@ import android.os.Bundle;
 import android.view.LayoutInflater;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Observer;
 
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor;
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.themedicon.ThemedIconSectionView;
 import com.android.wallpaper.R;
 import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WorkspaceViewModel;
 
 /** The {@link CustomizationSectionController} for themed icon section. */
 public class ThemedIconSectionController implements
@@ -33,16 +35,26 @@ public class ThemedIconSectionController implements
     private static final String KEY_THEMED_ICON_ENABLED = "SAVED_THEMED_ICON_ENABLED";
 
     private final ThemedIconSwitchProvider mThemedIconOptionsProvider;
-    private final WorkspaceViewModel mWorkspaceViewModel;
+    private final ThemedIconInteractor mInteractor;
+    private final ThemedIconSnapshotRestorer mSnapshotRestorer;
+    private final Observer<Boolean> mIsActivatedChangeObserver;
 
     private ThemedIconSectionView mThemedIconSectionView;
     private boolean mSavedThemedIconEnabled = false;
 
-
-    public ThemedIconSectionController(ThemedIconSwitchProvider themedIconOptionsProvider,
-            WorkspaceViewModel workspaceViewModel, @Nullable Bundle savedInstanceState) {
+    public ThemedIconSectionController(
+            ThemedIconSwitchProvider themedIconOptionsProvider,
+            ThemedIconInteractor interactor,
+            @Nullable Bundle savedInstanceState,
+            ThemedIconSnapshotRestorer snapshotRestorer) {
         mThemedIconOptionsProvider = themedIconOptionsProvider;
-        mWorkspaceViewModel = workspaceViewModel;
+        mInteractor = interactor;
+        mSnapshotRestorer = snapshotRestorer;
+        mIsActivatedChangeObserver = isActivated -> {
+            if (mThemedIconSectionView.isAttachedToWindow()) {
+                mThemedIconSectionView.getSwitch().setChecked(isActivated);
+            }
+        };
 
         if (savedInstanceState != null) {
             mSavedThemedIconEnabled = savedInstanceState.getBoolean(
@@ -64,15 +76,22 @@ public class ThemedIconSectionController implements
         mThemedIconSectionView.getSwitch().setChecked(mSavedThemedIconEnabled);
         mThemedIconOptionsProvider.fetchThemedIconEnabled(
                 enabled -> mThemedIconSectionView.getSwitch().setChecked(enabled));
+        mInteractor.isActivatedAsLiveData().observeForever(mIsActivatedChangeObserver);
         return mThemedIconSectionView;
     }
 
+    @Override
+    public void release() {
+        mInteractor.isActivatedAsLiveData().removeObserver(mIsActivatedChangeObserver);
+    }
+
     private void onViewActivated(Context context, boolean viewActivated) {
         if (context == null) {
             return;
         }
         mThemedIconOptionsProvider.setThemedIconEnabled(viewActivated);
-        mWorkspaceViewModel.getUpdateWorkspace().setValue(viewActivated);
+        mInteractor.setActivated(viewActivated);
+        mSnapshotRestorer.store(viewActivated);
     }
 
     @Override