Only go back when on complete apply

Test: Manully tested
Bug: 376091741
Flag: com.android.systemui.shared.new_customization_picker_ui
Change-Id: Ida9047f92568b0b9dcbb23061bf5d5e8e4c213f3
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 6282842..18b9584 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -64,12 +64,7 @@
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
                     viewModel.onApplyButtonClicked.collect { onApplyButtonClicked ->
-                        applyButton.setOnClickListener {
-                            if (onApplyButtonClicked != null) {
-                                onApplyButtonClicked.invoke()
-                                onNavBack.invoke()
-                            }
-                        }
+                        applyButton.setOnClickListener { onApplyButtonClicked?.invoke(onNavBack) }
                     }
                 }
 
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 2e500e2..b2e2039 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -131,7 +131,7 @@
         }
 
     @OptIn(ExperimentalCoroutinesApi::class)
-    val onApplyButtonClicked =
+    val onApplyButtonClicked: Flow<((onComplete: () -> Unit) -> Unit)?> =
         selectedOption
             .flatMapLatest {
                 when (it) {
@@ -155,8 +155,11 @@
             }
             .map { onApply ->
                 if (onApply != null) {
-                    fun() {
-                        viewModelScope.launch { onApply() }
+                    fun(onComplete: () -> Unit) {
+                        viewModelScope.launch {
+                            onApply()
+                            onComplete()
+                        }
                     }
                 } else {
                     null