diff --git a/src/com/android/customization/picker/mode/data/repository/DarkModeRepository.kt b/src/com/android/customization/picker/mode/data/repository/DarkModeRepository.kt
index 28f5017..69b489a 100644
--- a/src/com/android/customization/picker/mode/data/repository/DarkModeRepository.kt
+++ b/src/com/android/customization/picker/mode/data/repository/DarkModeRepository.kt
@@ -18,11 +18,9 @@
 
 import com.android.customization.picker.mode.shared.util.DarkModeUtil
 import com.android.wallpaper.system.PowerManagerWrapper
-import com.android.wallpaper.system.UiModeManagerWrapper
 import javax.inject.Inject
 import javax.inject.Singleton
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 
@@ -31,7 +29,7 @@
 @Inject
 constructor(
     darkModeUtil: DarkModeUtil,
-    private val uiModeManager: UiModeManagerWrapper,
+    private val darkModeStateRepository: DarkModeStateRepository,
     private val powerManager: PowerManagerWrapper,
 ) {
     private val isPowerSaveMode = MutableStateFlow(powerManager.getIsPowerSaveMode() ?: false)
@@ -43,17 +41,11 @@
             isPowerSaveMode.map { !it }
         } else flowOf(false)
 
-    private val _isDarkMode = MutableStateFlow(uiModeManager.getIsNightModeActivated())
-    val isDarkMode = _isDarkMode.asStateFlow()
+    val isDarkMode = darkModeStateRepository.isDarkMode
 
-    fun setDarkModeActivated(isActive: Boolean) {
-        uiModeManager.setNightModeActivated(isActive)
-        refreshIsDarkModeActivated()
-    }
+    fun setIsDarkMode(isActive: Boolean) = darkModeStateRepository.setIsDarkMode(isActive)
 
-    fun refreshIsDarkModeActivated() {
-        _isDarkMode.value = uiModeManager.getIsNightModeActivated()
-    }
+    fun refreshIsDarkMode() = darkModeStateRepository.refreshIsDarkMode()
 
     fun refreshIsPowerSaveModeActivated() {
         powerManager.getIsPowerSaveMode()?.let { isPowerSaveMode.value = it }
diff --git a/src/com/android/customization/picker/mode/domain/interactor/DarkModeInteractor.kt b/src/com/android/customization/picker/mode/domain/interactor/DarkModeInteractor.kt
index 1b74e33..e805222 100644
--- a/src/com/android/customization/picker/mode/domain/interactor/DarkModeInteractor.kt
+++ b/src/com/android/customization/picker/mode/domain/interactor/DarkModeInteractor.kt
@@ -25,5 +25,5 @@
     val isEnabled = repository.isEnabled
     val isDarkMode = repository.isDarkMode
 
-    fun setDarkModeActivated(isActive: Boolean) = repository.setDarkModeActivated(isActive)
+    fun setIsDarkMode(isActive: Boolean) = repository.setIsDarkMode(isActive)
 }
diff --git a/src/com/android/customization/picker/mode/shared/util/DarkModeLifecycleUtil.kt b/src/com/android/customization/picker/mode/shared/util/DarkModeLifecycleUtil.kt
index 749ac2e..f7bbfd2 100644
--- a/src/com/android/customization/picker/mode/shared/util/DarkModeLifecycleUtil.kt
+++ b/src/com/android/customization/picker/mode/shared/util/DarkModeLifecycleUtil.kt
@@ -59,7 +59,6 @@
             @Synchronized
             override fun onStart(owner: LifecycleOwner) {
                 super.onStart(owner)
-                darkModeRepository.refreshIsDarkModeActivated()
                 darkModeRepository.refreshIsPowerSaveModeActivated()
                 if (lifecycleOwner.lifecycle.currentState.isAtLeast(Lifecycle.State.STARTED)) {
                     activityContext.registerReceiver(
diff --git a/src/com/android/customization/picker/mode/ui/binder/DarkModeBinder.kt b/src/com/android/customization/picker/mode/ui/binder/DarkModeBinder.kt
index 9e9db7a..3e46288 100644
--- a/src/com/android/customization/picker/mode/ui/binder/DarkModeBinder.kt
+++ b/src/com/android/customization/picker/mode/ui/binder/DarkModeBinder.kt
@@ -54,7 +54,11 @@
                 }
                 launch {
                     viewModel.toggleDarkMode.collect {
-                        darkModeToggle.setOnCheckedChangeListener { _, _ -> it.invoke() }
+                        // Use onClickListener instead of onCheckedChangeListener to avoid the
+                        // potential cycle of: system value changes->the toggle isChecked value is
+                        // updated->the onCheckedChangeListener is called->the overriding value is
+                        // set. The overriding value should not be set by the system, only by user.
+                        darkModeToggle.setOnClickListener { _ -> it.invoke() }
                     }
                 }
             }
diff --git a/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt b/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
index f51d966..aacf001 100644
--- a/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
+++ b/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModel.kt
@@ -44,14 +44,17 @@
     val toggleDarkMode =
         combine(overridingIsDarkMode, isDarkMode) { override, current ->
             // Only set override if its value is different from current, else set to null
-            { _overridingIsDarkMode.value = if (override == null) !current else null }
+            {
+                _overridingIsDarkMode.value =
+                    if (override == null || override == current) !current else null
+            }
         }
 
     val onApply: Flow<(suspend () -> Unit)?> =
         combine(overridingIsDarkMode, isDarkMode, isEnabled) { override, current, isEnabled ->
             if (override != null && override != current && isEnabled) {
                 {
-                    interactor.setDarkModeActivated(override)
+                    interactor.setIsDarkMode(override)
                     logger.logDarkThemeApplied(override)
                 }
             } else null
diff --git a/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
index fbd56bc..3bdbb9d 100644
--- a/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/mode/ui/viewmodel/DarkModeViewModelTest.kt
@@ -49,7 +49,7 @@
     @Inject lateinit var darkModeRepository: DarkModeRepository
     @Inject lateinit var darkModeInteractor: DarkModeInteractor
     @Inject lateinit var logger: TestThemesUserEventLogger
-    lateinit var darkModeViewModel: DarkModeViewModel
+    private lateinit var darkModeViewModel: DarkModeViewModel
 
     @Inject lateinit var testDispatcher: TestDispatcher
     @Inject lateinit var testScope: TestScope
@@ -90,10 +90,11 @@
     fun toggleDarkMode() {
         testScope.runTest {
             uiModeManager.setNightModeActivated(false)
-            darkModeRepository.refreshIsDarkModeActivated()
+            darkModeRepository.refreshIsDarkMode()
             val getOverridingIsDarkMode = collectLastValue(darkModeViewModel.overridingIsDarkMode)
             val getPreviewingIsDarkMode = collectLastValue(darkModeViewModel.previewingIsDarkMode)
             val getToggleDarkMode = collectLastValue(darkModeViewModel.toggleDarkMode)
+            assertThat(getOverridingIsDarkMode()).isNull()
             assertThat(getPreviewingIsDarkMode()).isFalse()
 
             getToggleDarkMode()?.invoke()
@@ -109,10 +110,36 @@
     }
 
     @Test
+    fun previewingIsDarkMode_systemChanges() {
+        testScope.runTest {
+            uiModeManager.setNightModeActivated(false)
+            darkModeRepository.refreshIsDarkMode()
+            val getOverridingIsDarkMode = collectLastValue(darkModeViewModel.overridingIsDarkMode)
+            val getPreviewingIsDarkMode = collectLastValue(darkModeViewModel.previewingIsDarkMode)
+            assertThat(getOverridingIsDarkMode()).isNull()
+            assertThat(getPreviewingIsDarkMode()).isFalse()
+
+            // Turn on dark mode
+            uiModeManager.setNightModeActivated(true)
+            darkModeRepository.refreshIsDarkMode()
+
+            assertThat(getOverridingIsDarkMode()).isNull()
+            assertThat(getPreviewingIsDarkMode()).isTrue()
+
+            // Turn off dark mode
+            uiModeManager.setNightModeActivated(false)
+            darkModeRepository.refreshIsDarkMode()
+
+            assertThat(getOverridingIsDarkMode()).isNull()
+            assertThat(getPreviewingIsDarkMode()).isFalse()
+        }
+    }
+
+    @Test
     fun onApply_shouldLogDarkTheme() {
         testScope.runTest {
             uiModeManager.setNightModeActivated(false)
-            darkModeRepository.refreshIsDarkModeActivated()
+            darkModeRepository.refreshIsDarkMode()
             val getToggleDarkMode = collectLastValue(darkModeViewModel.toggleDarkMode)
             val onApply = collectLastValue(darkModeViewModel.onApply)
 
@@ -127,7 +154,7 @@
     fun onApply_shouldApplyDarkTheme() {
         testScope.runTest {
             uiModeManager.setNightModeActivated(false)
-            darkModeRepository.refreshIsDarkModeActivated()
+            darkModeRepository.refreshIsDarkMode()
             val getToggleDarkMode = collectLastValue(darkModeViewModel.toggleDarkMode)
             val onApply = collectLastValue(darkModeViewModel.onApply)
 
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index 649c298..a8edc4d 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -26,15 +26,18 @@
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor2
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
+import com.android.customization.picker.mode.data.repository.DarkModeStateRepository
 import com.android.systemui.monet.Style
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
-import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import dagger.hilt.android.internal.lifecycle.RetainedLifecycleImpl
+import dagger.hilt.android.testing.HiltAndroidRule
+import dagger.hilt.android.testing.HiltAndroidTest
+import javax.inject.Inject
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.launch
@@ -45,36 +48,40 @@
 import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
 
+@HiltAndroidTest
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(RobolectricTestRunner::class)
 class ColorPickerViewModel2Test {
+    @get:Rule var hiltRule = HiltAndroidRule(this)
+
     private val logger = TestThemesUserEventLogger()
     private lateinit var underTest: ColorPickerViewModel2
-    private lateinit var repository: FakeColorPickerRepository2
-    private lateinit var interactor: ColorPickerInteractor2
-    private lateinit var store: FakeSnapshotStore
     private lateinit var colorUpdateViewModel: ColorUpdateViewModel
 
     private lateinit var context: Context
     private lateinit var testScope: TestScope
 
+    @Inject lateinit var repository: FakeColorPickerRepository2
+    @Inject lateinit var interactor: ColorPickerInteractor2
+    @Inject lateinit var darkModeStateRepository: DarkModeStateRepository
+
     @Before
     fun setUp() {
+        hiltRule.inject()
+
         context = InstrumentationRegistry.getInstrumentation().targetContext
         val testDispatcher = UnconfinedTestDispatcher()
         Dispatchers.setMain(testDispatcher)
         testScope = TestScope(testDispatcher)
-        repository = FakeColorPickerRepository2()
-        store = FakeSnapshotStore()
 
-        interactor = ColorPickerInteractor2(repository = repository)
-
-        colorUpdateViewModel = ColorUpdateViewModel(context, RetainedLifecycleImpl())
+        colorUpdateViewModel =
+            ColorUpdateViewModel(context, RetainedLifecycleImpl(), darkModeStateRepository)
 
         underTest =
             ColorPickerViewModel2(