Reset support for dark theme.

Fix: 267804364
Test: unit tests added
Test: manually verified that toggling the dark theme button makes the
RESET button show up, toggling it back, make it disappear. If the RESET
button is clicked, the theme setting goes back to the original.

Change-Id: Ifad8023928f20718a798b24d7da77698bc9d1677
diff --git a/src/com/android/customization/model/mode/DarkModeSectionController.java b/src/com/android/customization/model/mode/DarkModeSectionController.java
index f56b709..ebeaa56 100644
--- a/src/com/android/customization/model/mode/DarkModeSectionController.java
+++ b/src/com/android/customization/model/mode/DarkModeSectionController.java
@@ -59,12 +59,17 @@
 
     private Context mContext;
     private DarkModeSectionView mDarkModeSectionView;
+    private final DarkModeSnapshotRestorer mSnapshotRestorer;
 
-    public DarkModeSectionController(Context context, Lifecycle lifecycle) {
+    public DarkModeSectionController(
+            Context context,
+            Lifecycle lifecycle,
+            DarkModeSnapshotRestorer snapshotRestorer) {
         mContext = context;
         mLifecycle = lifecycle;
         mPowerManager = context.getSystemService(PowerManager.class);
         mLifecycle.addObserver(this);
+        mSnapshotRestorer = snapshotRestorer;
     }
 
     @OnLifecycleEvent(Lifecycle.Event.ON_START)
@@ -132,6 +137,7 @@
                     mDarkModeSectionView.announceForAccessibility(
                             context.getString(R.string.mode_changed));
                     uiModeManager.setNightModeActivated(viewActivated);
+                    mSnapshotRestorer.store(viewActivated);
                 },
                 /* delayMillis= */ shortDelay);
     }
diff --git a/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
new file mode 100644
index 0000000..aa8d97c
--- /dev/null
+++ b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2023 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.customization.model.mode
+
+import android.app.UiModeManager
+import android.content.Context
+import android.content.res.Configuration
+import androidx.annotation.VisibleForTesting
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+class DarkModeSnapshotRestorer : SnapshotRestorer {
+
+    private val backgroundDispatcher: CoroutineDispatcher
+    private val isActive: () -> Boolean
+    private val setActive: suspend (Boolean) -> Unit
+
+    private lateinit var store: SnapshotStore
+
+    constructor(
+        context: Context,
+        manager: UiModeManager,
+        backgroundDispatcher: CoroutineDispatcher,
+    ) : this(
+        backgroundDispatcher = backgroundDispatcher,
+        isActive = {
+            context.applicationContext.resources.configuration.uiMode and
+                Configuration.UI_MODE_NIGHT_YES != 0
+        },
+        setActive = { isActive -> manager.setNightModeActivated(isActive) },
+    )
+
+    @VisibleForTesting
+    constructor(
+        backgroundDispatcher: CoroutineDispatcher,
+        isActive: () -> Boolean,
+        setActive: suspend (Boolean) -> Unit,
+    ) {
+        this.backgroundDispatcher = backgroundDispatcher
+        this.isActive = isActive
+        this.setActive = setActive
+    }
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        return snapshot(
+            isActivated = isActive(),
+        )
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isActivated = snapshot.args[KEY]?.toBoolean() == true
+        withContext(backgroundDispatcher) { setActive(isActivated) }
+    }
+
+    fun store(
+        isActivated: Boolean,
+    ) {
+        store.store(
+            snapshot(
+                isActivated = isActivated,
+            ),
+        )
+    }
+
+    private fun snapshot(
+        isActivated: Boolean,
+    ): RestorableSnapshot {
+        return RestorableSnapshot(
+            args =
+                buildMap {
+                    put(
+                        KEY,
+                        isActivated.toString(),
+                    )
+                }
+        )
+    }
+
+    companion object {
+        private const val KEY = "is_activated"
+    }
+}
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 16d4bc7..d05e459 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -12,6 +12,7 @@
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.GridSectionController;
 import com.android.customization.model.mode.DarkModeSectionController;
+import com.android.customization.model.mode.DarkModeSnapshotRestorer;
 import com.android.customization.model.themedicon.ThemedIconSectionController;
 import com.android.customization.model.themedicon.ThemedIconSwitchProvider;
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider;
@@ -56,6 +57,7 @@
             mClockCarouselViewModelProvider;
     private final PreviewWithClockCarouselSectionController.ClockViewFactoryProvider
             mClockViewFactoryProvider;
+    private final DarkModeSnapshotRestorer mDarkModeSnapshotRestorer;
 
     public DefaultCustomizationSections(
             KeyguardQuickAffordancePickerInteractor keyguardQuickAffordancePickerInteractor,
@@ -65,7 +67,8 @@
             BaseFlags flags,
             ClockRegistryProvider clockRegistryProvider,
             ClockCarouselViewModelProvider clockCarouselViewModelProvider,
-            ClockViewFactoryProvider clockViewFactoryProvider) {
+            ClockViewFactoryProvider clockViewFactoryProvider,
+            DarkModeSnapshotRestorer darkModeSnapshotRestorer) {
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
                 keyguardQuickAffordancePickerViewModelFactory;
@@ -74,6 +77,7 @@
         mClockRegistryProvider = clockRegistryProvider;
         mClockCarouselViewModelProvider = clockCarouselViewModelProvider;
         mClockViewFactoryProvider = clockViewFactoryProvider;
+        mDarkModeSnapshotRestorer = darkModeSnapshotRestorer;
     }
 
     @Override
@@ -157,8 +161,10 @@
 
             case HOME_SCREEN:
                 // Dark/Light theme section.
-                sectionControllers.add(new DarkModeSectionController(activity,
-                        lifecycleOwner.getLifecycle()));
+                sectionControllers.add(new DarkModeSectionController(
+                        activity,
+                        lifecycleOwner.getLifecycle(),
+                        mDarkModeSnapshotRestorer));
 
                 // Themed app icon section.
                 sectionControllers.add(new ThemedIconSectionController(
@@ -198,8 +204,10 @@
                 activity, wallpaperColorsViewModel, lifecycleOwner, savedInstanceState));
 
         // Dark/Light theme section.
-        sectionControllers.add(new DarkModeSectionController(activity,
-                lifecycleOwner.getLifecycle()));
+        sectionControllers.add(new DarkModeSectionController(
+                activity,
+                lifecycleOwner.getLifecycle(),
+                mDarkModeSnapshotRestorer));
 
         // Themed app icon section.
         sectionControllers.add(new ThemedIconSectionController(
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index c2a6565..007b419 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -15,6 +15,7 @@
  */
 package com.android.customization.module
 
+import android.app.UiModeManager
 import android.content.Context
 import android.content.Intent
 import android.net.Uri
@@ -22,6 +23,7 @@
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
+import com.android.customization.model.mode.DarkModeSnapshotRestorer
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
@@ -86,6 +88,7 @@
     private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
     private var colorPickerInteractor: ColorPickerInteractor? = null
     private var colorPickerViewModelFactory: ColorPickerViewModel.Factory? = null
+    private var darkModeSnapshotRestorer: DarkModeSnapshotRestorer? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
@@ -112,7 +115,8 @@
                                 registry = registry,
                             )
                         }
-                    }
+                    },
+                    getDarkModeSnapshotRestorer(activity),
                 )
                 .also { customizationSections = it }
     }
@@ -173,6 +177,7 @@
                 getKeyguardQuickAffordanceSnapshotRestorer(context)
             this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
             this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] = getNotificationsSnapshotRestorer(context)
+            this[KEY_DARK_MODE_SNAPSHOT_RESTORER] = getDarkModeSnapshotRestorer(context)
         }
     }
 
@@ -348,6 +353,18 @@
                 .also { colorPickerViewModelFactory = it }
     }
 
+    protected fun getDarkModeSnapshotRestorer(
+        context: Context,
+    ): DarkModeSnapshotRestorer {
+        return darkModeSnapshotRestorer
+            ?: DarkModeSnapshotRestorer(
+                    context = context,
+                    manager = context.getSystemService(Context.UI_MODE_SERVICE) as UiModeManager,
+                    backgroundDispatcher = Dispatchers.IO,
+                )
+                .also { darkModeSnapshotRestorer = it }
+    }
+
     companion object {
         @JvmStatic
         private val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
@@ -356,11 +373,13 @@
         private val KEY_WALLPAPER_SNAPSHOT_RESTORER = KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1
         @JvmStatic
         private val KEY_NOTIFICATIONS_SNAPSHOT_RESTORER = KEY_WALLPAPER_SNAPSHOT_RESTORER + 1
+        @JvmStatic
+        private val KEY_DARK_MODE_SNAPSHOT_RESTORER = KEY_NOTIFICATIONS_SNAPSHOT_RESTORER + 1
 
         /**
          * When this injector is overridden, this is the minimal value that should be used by
          * restorers returns in [getSnapshotRestorers].
          */
-        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_NOTIFICATIONS_SNAPSHOT_RESTORER + 1
+        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_DARK_MODE_SNAPSHOT_RESTORER + 1
     }
 }
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
new file mode 100644
index 0000000..38067b7
--- /dev/null
+++ b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2023 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.customization.model.mode
+
+import androidx.test.filters.SmallTest
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class DarkModeSnapshotRestorerTest {
+
+    private lateinit var underTest: DarkModeSnapshotRestorer
+    private lateinit var testScope: TestScope
+
+    private var isActive = false
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        underTest =
+            DarkModeSnapshotRestorer(
+                backgroundDispatcher = testDispatcher,
+                isActive = { isActive },
+                setActive = { isActive = it },
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() =
+        testScope.runTest {
+            isActive = true
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up and restore - inactive`() =
+        testScope.runTest {
+            isActive = false
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isFalse()
+        }
+
+    @Test
+    fun `set up - deactivate - restore to active`() =
+        testScope.runTest {
+            isActive = true
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = false)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up - activate - restore to inactive`() =
+        testScope.runTest {
+            isActive = false
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = true)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isFalse()
+        }
+}