Merge "Revert "Clean up and refactor load initial colors (1/3)"" into udc-dev

10 files changed, 40 insertions, 62 deletions

diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 00f9433c..b408a89d 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -145,7 +145,6 @@ public final class DefaultCustomizationSections implements CustomizationSections
                                         activity,
                                         mColorPickerViewModelFactory)
                                         .get(ColorPickerViewModel.class),
-                                wallpaperColorsViewModel,
                                 lifecycleOwner),
                         // Wallpaper quick switch section.
                         new WallpaperQuickSwitchSectionController(
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 1e0dbe5f..d00ed28e 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -439,8 +439,7 @@ open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInject
                     repository =
                         ColorPickerRepositoryImpl(
                             wallpaperColorsViewModel,
-                            getColorCustomizationManager(appContext),
-                            WallpaperManager.getInstance(appContext),
+                            getColorCustomizationManager(appContext)
                         ),
                     snapshotRestorer = {
                         getColorPickerSnapshotRestorer(appContext, wallpaperColorsViewModel)
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
index ba77f25f..7cf9fd03 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -37,7 +37,4 @@ interface ColorPickerRepository {
 
     /** Returns the current selected color source based on system settings */
     fun getCurrentColorSource(): String?
-
-    /** Retrieves and stores the wallpaper colors for generating wallpaper color options */
-    suspend fun loadInitialColors()
 }
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index f1c695dd..41ef3a57 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -16,7 +16,6 @@
  */
 package com.android.customization.picker.color.data.repository
 
-import android.app.WallpaperManager
 import android.util.Log
 import com.android.customization.model.CustomizationManager
 import com.android.customization.model.color.ColorCustomizationManager
@@ -27,20 +26,17 @@ import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
 import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
-import kotlinx.coroutines.withContext
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
 // TODO (b/268203200): Create test for ColorPickerRepositoryImpl
 class ColorPickerRepositoryImpl(
-    private val wallpaperColorsViewModel: WallpaperColorsViewModel,
+    wallpaperColorsViewModel: WallpaperColorsViewModel,
     private val colorManager: ColorCustomizationManager,
-    private val wallpaperManager: WallpaperManager,
 ) : ColorPickerRepository {
 
     private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
@@ -155,17 +151,6 @@ class ColorPickerRepositoryImpl(
         return colorManager.currentColorSource
     }
 
-    override suspend fun loadInitialColors() {
-        withContext(Dispatchers.IO) {
-            val lockColors = wallpaperManager.getWallpaperColors(WallpaperManager.FLAG_LOCK)
-            val homeColors = wallpaperManager.getWallpaperColors(WallpaperManager.FLAG_SYSTEM)
-            withContext(Dispatchers.Main) {
-                wallpaperColorsViewModel.setLockWallpaperColors(lockColors)
-                wallpaperColorsViewModel.setHomeWallpaperColors(homeColors)
-            }
-        }
-    }
-
     private fun ColorOptionImpl.toModel(): ColorOptionModel {
         return ColorOptionModel(
             key = "${this.type}::${this.style}::${this.serializedPackages}",
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
index e4b8795e..714129df 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -166,8 +166,6 @@ class FakeColorPickerRepository(private val context: Context) : ColorPickerRepos
             else -> null
         }
 
-    override suspend fun loadInitialColors() {}
-
     private fun ColorOptionModel.testEquals(other: Any?): Boolean {
         if (other == null) {
             return false
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
index 36d2a0d4..8c7a4b72 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -46,6 +46,4 @@ class ColorPickerInteractor(
     }
 
     fun getCurrentColorOption(): ColorOptionModel = repository.getCurrentColorOption()
-
-    suspend fun loadInitialColors() = repository.loadInitialColors()
 }
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index 3390a0d1..78bfa43e 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -28,7 +28,6 @@ import androidx.lifecycle.lifecycleScope
 import com.android.customization.model.mode.DarkModeSectionController
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.color.ui.binder.ColorPickerBinder
-import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -74,29 +73,20 @@ class ColorPickerFragment : AppbarFragment() {
         val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
         val displayUtils: DisplayUtils = injector.getDisplayUtils(requireContext())
         val wcViewModel = injector.getWallpaperColorsViewModel()
-
-        val colorPickerViewModel =
-            ViewModelProvider(
-                    requireActivity(),
-                    injector.getColorPickerViewModelFactory(
-                        context = requireContext(),
-                        wallpaperColorsViewModel = wcViewModel,
-                    ),
-                )
-                .get() as ColorPickerViewModel
-
-        // load wallpaper colors if it has not been populated
-        if (
-            wcViewModel.lockWallpaperColors.value is WallpaperColorsModel.Loading ||
-                wcViewModel.homeWallpaperColors.value is WallpaperColorsModel.Loading
-        ) {
-            lifecycleScope.launch { colorPickerViewModel.loadInitialColors() }
-        }
+        val wallpaperManager = WallpaperManager.getInstance(requireContext())
 
         binding =
             ColorPickerBinder.bind(
                 view = view,
-                viewModel = colorPickerViewModel,
+                viewModel =
+                    ViewModelProvider(
+                            requireActivity(),
+                            injector.getColorPickerViewModelFactory(
+                                context = requireContext(),
+                                wallpaperColorsViewModel = wcViewModel,
+                            ),
+                        )
+                        .get(),
                 lifecycleOwner = this,
             )
 
@@ -120,6 +110,18 @@ class ColorPickerFragment : AppbarFragment() {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
                                 { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wcViewModel.lockWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wcViewModel,
+                                                CustomizationSections.Screen.LOCK_SCREEN
+                                            )
+                                        }
+                                    }
                                     continuation.resume(lockWallpaper ?: homeWallpaper, null)
                                 },
                                 forceReload,
@@ -155,6 +157,18 @@ class ColorPickerFragment : AppbarFragment() {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
                                 { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wcViewModel.homeWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wcViewModel,
+                                                CustomizationSections.Screen.HOME_SCREEN
+                                            )
+                                        }
+                                    }
                                     continuation.resume(homeWallpaper ?: lockWallpaper, null)
                                 },
                                 forceReload,
diff --git a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
index d0ddb2b0..f1c982b4 100644
--- a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
+++ b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
@@ -20,7 +20,6 @@ package com.android.customization.picker.color.ui.section
 import android.content.Context
 import android.view.LayoutInflater
 import androidx.lifecycle.LifecycleOwner
-import androidx.lifecycle.lifecycleScope
 import com.android.customization.picker.color.ui.binder.ColorSectionViewBinder
 import com.android.customization.picker.color.ui.fragment.ColorPickerFragment
 import com.android.customization.picker.color.ui.view.ColorSectionView2
@@ -28,14 +27,10 @@ import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.CustomizationSectionController
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController as NavigationController
-import com.android.wallpaper.model.WallpaperColorsModel
-import com.android.wallpaper.model.WallpaperColorsViewModel
-import kotlinx.coroutines.launch
 
 class ColorSectionController2(
     private val navigationController: NavigationController,
     private val viewModel: ColorPickerViewModel,
-    private val wcViewModel: WallpaperColorsViewModel,
     private val lifecycleOwner: LifecycleOwner
 ) : CustomizationSectionController<ColorSectionView2> {
 
@@ -58,15 +53,6 @@ class ColorSectionController2(
                     R.layout.color_section_view2,
                     null,
                 ) as ColorSectionView2
-
-        // load wallpaper colors if it has not been populated
-        if (
-            wcViewModel.lockWallpaperColors.value is WallpaperColorsModel.Loading ||
-                wcViewModel.homeWallpaperColors.value is WallpaperColorsModel.Loading
-        ) {
-            lifecycleOwner.lifecycleScope.launch { viewModel.loadInitialColors() }
-        }
-
         ColorSectionViewBinder.bind(
             view = view,
             viewModel = viewModel,
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 057eff5d..67c68387 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -202,8 +202,6 @@ private constructor(
                 replay = 1,
             )
 
-    suspend fun loadInitialColors() = interactor.loadInitialColors()
-
     class Factory(
         private val context: Context,
         private val interactor: ColorPickerInteractor,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index b6301225..56c6c30a 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -94,6 +94,10 @@ open class PreviewWithThemeSectionController(
                                 } else {
                                     homeWallpaper ?: lockWallpaper
                                 }
+                            loadInitialColors(
+                                context = context,
+                                screen = screen,
+                            )
                             continuation.resume(wallpaper, null)
                         },
                         forceReload,