Add dark theme custom type into the dark theme settings page

Also fix some test cases.

I'd prefer using Robolectric shadows rather than mocking Android
framework classes. However, I can't do that in this CL due to the
lack of multiple shadow classes support.

Test: make RunSettingsRoboTests -j56
      ROBOTEST_FILTER=com.android.settings.display.darkmode.*
Bug: 215182463
Change-Id: If557d2933927a5dd0fadd3f6db6bb6f0ab7dd5ee

tests/robotests/src/com/android/settings/display/darkmode/DarkModeActivationPreferenceControllerTest.java

@@ -19,7 +19,6 @@ import static com.android.settings.core.BasePreferenceController.AVAILABLE_UNSEA
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -54,12 +53,11 @@ public class DarkModeActivationPreferenceControllerTest {
     private DarkModeActivationPreferenceController mController;
     private String mPreferenceKey = "key";
 
-    @Mock
     private MainSwitchPreference mPreference;
     @Mock
     private PreferenceScreen mScreen;
     @Mock
-    private Resources res;
+    private Resources mRes;
     @Mock
     private UiModeManager mService;
     @Mock
@@ -77,9 +75,10 @@ public class DarkModeActivationPreferenceControllerTest {
         MockitoAnnotations.initMocks(this);
         FakeFeatureFactory.setupForTest();
         mContext = spy(RuntimeEnvironment.application);
+        mPreference = new MainSwitchPreference(mContext);
         mService = mock(UiModeManager.class);
-        when(mContext.getResources()).thenReturn(res);
-        when(res.getConfiguration()).thenReturn(mConfigNightNo);
+        when(mContext.getResources()).thenReturn(mRes);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightNo);
         when(mContext.getSystemService(UiModeManager.class)).thenReturn(mService);
         when(mContext.getSystemService(PowerManager.class)).thenReturn(mPM);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
@@ -101,10 +100,14 @@ public class DarkModeActivationPreferenceControllerTest {
                 R.string.dark_ui_summary_off_auto_mode_never)).thenReturn("summary_off_manual");
         when(mContext.getString(
                 R.string.dark_ui_summary_on_auto_mode_never)).thenReturn("summary_on_manual");
-        when(mContext.getString(
-                R.string.dark_ui_summary_on_auto_mode_custom)).thenReturn("summary_on_custom");
-        when(mContext.getString(
-                R.string.dark_ui_summary_off_auto_mode_custom)).thenReturn("summary_off_custom");
+        when(mContext.getString(R.string.dark_ui_summary_on_auto_mode_custom, "10:00 AM"))
+                .thenReturn("summary_on_custom");
+        when(mContext.getString(R.string.dark_ui_summary_off_auto_mode_custom, "10:00 AM"))
+                .thenReturn("summary_off_custom");
+        when(mContext.getString(R.string.dark_ui_summary_on_auto_mode_custom_bedtime))
+                .thenReturn("summary_on_custom_bedtime");
+        when(mContext.getString(R.string.dark_ui_summary_off_auto_mode_custom_bedtime))
+                .thenReturn("summary_off_custom_bedtime");
         mController = new DarkModeActivationPreferenceController(mContext, mPreferenceKey, mFormat);
         mController.displayPreference(mScreen);
         mConfigNightNo.uiMode = Configuration.UI_MODE_NIGHT_NO;
@@ -114,85 +117,95 @@ public class DarkModeActivationPreferenceControllerTest {
     }
 
     @Test
-    public void nightMode_toggleButton_offManual() {
+    public void nightMode_toggleButton_onManual() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_YES);
-        when(res.getConfiguration()).thenReturn(mConfigNightYes);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightYes);
 
         mController.updateState(mPreference);
 
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_on_manual");
     }
 
     @Test
-    public void nightMode_toggleButton_offCustom() {
-        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
-        when(res.getConfiguration()).thenReturn(mConfigNightYes);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
+    public void nightMode_toggleButton_offManual() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_NO);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightNo);
 
-        mController.updateState(preference);
+        mController.updateState(mPreference);
 
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_off_manual");
     }
 
     @Test
     public void nightMode_toggleButton_onCustom() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
-        when(res.getConfiguration()).thenReturn(mConfigNightYes);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightYes);
 
-        mController.updateState(preference);
+        mController.updateState(mPreference);
 
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_on_custom");
     }
 
     @Test
-    public void nightMode_toggleButton_onAutoWhenModeIsYes() {
-        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_YES);
-        when(res.getConfiguration()).thenReturn(mConfigNightNo);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext, null);
+    public void nightMode_toggleButton_offCustom() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightNo);
 
-        mController.updateState(preference);
+        mController.updateState(mPreference);
 
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_off_custom");
     }
 
     @Test
-    public void nightMode_toggleButton_onAutoWhenModeIsAuto() {
+    public void nightMode_toggleButton_onCustomBedtime() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightYes);
+
+        mController.updateState(mPreference);
+
+        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_on_custom_bedtime");
+    }
+
+    @Test
+    public void nightMode_toggleButton_offCustomBedtime() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightNo);
+
+        mController.updateState(mPreference);
+
+        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_off_custom_bedtime");
+    }
+
+    @Test
+    public void nightMode_toggleButton_onAuto() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_AUTO);
-        when(res.getConfiguration()).thenReturn(mConfigNightNo);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightYes);
 
-        mController.updateState(preference);
+        mController.updateState(mPreference);
 
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isTrue();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_on_auto");
     }
 
     @Test
-    public void nightModeSummary_buttonText_onManual() {
-        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_NO);
-        when(res.getConfiguration()).thenReturn(mConfigNightYes);
-
-        assertEquals(mController.getSummary(), mContext.getString(
-                R.string.dark_ui_summary_on_auto_mode_never));
-    }
-
-    @Test
-    public void nightModeSummary_buttonText_offAuto() {
+    public void nightMode_toggleButton_offAuto() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_AUTO);
-        when(res.getConfiguration()).thenReturn(mConfigNightNo);
+        when(mRes.getConfiguration()).thenReturn(mConfigNightNo);
 
-        assertEquals(mController.getSummary(), mContext.getString(
-                R.string.dark_ui_summary_off_auto_mode_auto));
-    }
+        mController.updateState(mPreference);
 
-    @Test
-    public void buttonVisisbility_hideButton_offWhenInPowerSaveMode() {
-        when(mPM.isPowerSaveMode()).thenReturn(true);
-        final MainSwitchPreference preference = new MainSwitchPreference(mContext);
-
-        mController.updateState(preference);
-        assertThat(preference.isChecked()).isFalse();
+        assertThat(mPreference.isChecked()).isFalse();
+        assertThat(mController.getSummary().toString()).isEqualTo("summary_off_auto");
     }
 
     @Test

tests/robotests/src/com/android/settings/display/darkmode/DarkModeCustomBedtimePreferenceControllerTest.java

@@ -0,0 +1,161 @@
+/*
+ * Copyright (C) 2022 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.settings.display.darkmode;
+
+import static android.app.UiModeManager.MODE_NIGHT_AUTO;
+import static android.app.UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME;
+import static android.app.UiModeManager.MODE_NIGHT_CUSTOM_TYPE_SCHEDULE;
+import static android.app.UiModeManager.MODE_NIGHT_NO;
+import static android.app.UiModeManager.MODE_NIGHT_YES;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE_UNSEARCHABLE;
+import static com.android.settings.core.BasePreferenceController.UNSUPPORTED_ON_DEVICE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.UiModeManager;
+import android.content.Context;
+import android.content.res.Resources;
+
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settings.testutils.BedtimeSettingsUtils;
+import com.android.settings.testutils.FakeFeatureFactory;
+import com.android.settingslib.widget.FooterPreference;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+
+@RunWith(RobolectricTestRunner.class)
+public class DarkModeCustomBedtimePreferenceControllerTest {
+    @Mock
+    private UiModeManager mService;
+    @Mock
+    private PreferenceScreen mScreen;
+    @Mock
+    private Resources mResources;
+    @Mock
+    private FooterPreference mFooterPreference;
+
+    private DarkModeCustomBedtimePreferenceController mController;
+    private Context mContext;
+    private BedtimeSettingsUtils mBedtimeSettingsUtils;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        FakeFeatureFactory.setupForTest();
+
+        mContext = spy(ApplicationProvider.getApplicationContext());
+        mBedtimeSettingsUtils = new BedtimeSettingsUtils(mContext);
+
+        when(mContext.getSystemService(UiModeManager.class)).thenReturn(mService);
+
+        when(mContext.getResources()).thenReturn(mResources);
+        when(mResources.getString(com.android.internal.R.string.config_defaultWellbeingPackage))
+                .thenReturn("wellbeing");
+
+        when(mScreen.findPreference(anyString())).thenReturn(mFooterPreference);
+
+        mController = new DarkModeCustomBedtimePreferenceController(mContext, "key");
+    }
+
+    @Test
+    public void getAvailabilityStatus_bedtimeSettingsExist_shouldBeAvailableUnsearchable() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightModeCustomType()).thenReturn(MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE_UNSEARCHABLE);
+    }
+
+    @Test
+    public void getAvailabilityStatus_bedtimeSettingsDisabled_shouldBeUnsupportedOnDevice() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                false /* enabled */);
+        when(mService.getNightModeCustomType()).thenReturn(MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(UNSUPPORTED_ON_DEVICE);
+    }
+
+    @Test
+    public void nightModeCustomModeBedtime_bedtimeSettingsExist_shouldShowFooterPreference() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightModeCustomType()).thenReturn(MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+
+        mController.updateState(mFooterPreference);
+
+        verify(mFooterPreference).setVisible(eq(true));
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE_UNSEARCHABLE);
+    }
+
+    @Test
+    public void nightModeCustomModeSchedule_bedtimeSettingsExist_shouldHideFooterPreference() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightModeCustomType()).thenReturn(MODE_NIGHT_CUSTOM_TYPE_SCHEDULE);
+
+        mController.updateState(mFooterPreference);
+
+        verify(mFooterPreference).setVisible(eq(false));
+    }
+
+    @Test
+    public void nightModeNo_bedtimeSettingsExist_shouldHideFooterPreference() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightMode()).thenReturn(MODE_NIGHT_NO);
+
+        mController.updateState(mFooterPreference);
+
+        verify(mFooterPreference).setVisible(eq(false));
+    }
+
+    @Test
+    public void nightModeYes_bedtimeSettingsExist_shouldHideFooterPreference() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightMode()).thenReturn(MODE_NIGHT_YES);
+
+        mController.updateState(mFooterPreference);
+
+        verify(mFooterPreference).setVisible(eq(false));
+    }
+
+    @Test
+    public void nightModeAuto_bedtimeSettingsExist_shouldHideFooterPreference() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mService.getNightMode()).thenReturn(MODE_NIGHT_AUTO);
+
+        mController.updateState(mFooterPreference);
+
+        verify(mFooterPreference).setVisible(eq(false));
+    }
+}

tests/robotests/src/com/android/settings/display/darkmode/DarkModeCustomPreferenceControllerTest.java

@@ -15,12 +15,20 @@
 
 package com.android.settings.display.darkmode;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 import android.app.UiModeManager;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.content.res.Configuration;
 import android.content.res.Resources;
+
 import androidx.preference.Preference;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -28,12 +36,6 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
 @RunWith(RobolectricTestRunner.class)
 public class DarkModeCustomPreferenceControllerTest {
     private DarkModeCustomPreferenceController mController;
@@ -69,23 +71,87 @@ public class DarkModeCustomPreferenceControllerTest {
     }
 
     @Test
-    public void nightMode_customOff_hidePreference() {
+    public void nightMode_manualOn_hidePreference() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_YES);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_YES;
+
         mController.refreshSummary(mPreference);
+
         verify(mPreference).setVisible(eq(false));
     }
 
+    @Test
+    public void nightMode_manualOff_hidePreference() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_NO);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_NO;
+
+        mController.refreshSummary(mPreference);
+
+        verify(mPreference).setVisible(eq(false));
+    }
+
+
+    @Test
+    public void nightMode_customOn_showPreference() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_YES;
+
+        mController.refreshSummary(mPreference);
+
+        verify(mPreference).setVisible(eq(true));
+    }
+
     @Test
     public void nightMode_customOff_showPreference() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_NO;
+
         mController.refreshSummary(mPreference);
+
         verify(mPreference).setVisible(eq(true));
     }
 
     @Test
-    public void nightMode_customOff_setSummaryNotNull() {
+    public void nightMode_customBedtimeOn_hidePreference() {
         when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_YES;
+
         mController.refreshSummary(mPreference);
-        verify(mPreference).setSummary(eq(mFormat.of(null)));
+
+        verify(mPreference).setVisible(eq(false));
+    }
+
+    @Test
+    public void nightMode_customBedtimeOff_hidePreference() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_NO;
+
+        mController.refreshSummary(mPreference);
+
+        verify(mPreference).setVisible(eq(false));
+    }
+
+    @Test
+    public void nightMode_customOn_setSummaryTo10Am() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_YES;
+
+        mController.refreshSummary(mPreference);
+
+        verify(mPreference).setSummary(eq("10:00 AM"));
+    }
+
+    @Test
+    public void nightMode_customOff_setSummaryTo10Am() {
+        when(mService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        mConfig.uiMode = Configuration.UI_MODE_NIGHT_NO;
+
+        mController.refreshSummary(mPreference);
+
+        verify(mPreference).setSummary(eq("10:00 AM"));
     }
 }

tests/robotests/src/com/android/settings/display/darkmode/DarkModeObserverTest.java

@@ -14,35 +14,58 @@
 
 package com.android.settings.display.darkmode;
 
-import android.content.Context;
-import android.database.ContentObserver;
-import android.net.Uri;
-import android.provider.Settings;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mockito;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
+import static android.provider.Settings.Secure.DARK_THEME_CUSTOM_END_TIME;
+import static android.provider.Settings.Secure.DARK_THEME_CUSTOM_START_TIME;
+import static android.provider.Settings.Secure.UI_NIGHT_MODE;
+import static android.provider.Settings.Secure.UI_NIGHT_MODE_CUSTOM_TYPE;
+import static android.provider.Settings.Secure.getUriFor;
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 
+import android.content.Context;
+import android.database.ContentObserver;
+import android.net.Uri;
 
-@RunWith(RobolectricTestRunner.class)
+import androidx.test.core.app.ApplicationProvider;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.ParameterizedRobolectricTestRunner;
+
+import java.util.Arrays;
+import java.util.List;
+
+@RunWith(ParameterizedRobolectricTestRunner.class)
 public class DarkModeObserverTest {
+    @ParameterizedRobolectricTestRunner.Parameters(name = "uri: {0}")
+    public static List<Object[]> params() {
+        return Arrays.asList(
+                new Object[]{getUriFor(UI_NIGHT_MODE)},
+                new Object[]{getUriFor(UI_NIGHT_MODE_CUSTOM_TYPE)},
+                new Object[]{getUriFor(DARK_THEME_CUSTOM_START_TIME)},
+                new Object[]{getUriFor(DARK_THEME_CUSTOM_END_TIME)});
+    }
+
+    private final Uri mUri;
+
     private Context mContext;
     private ContentObserver mContentObserver;
     private DarkModeObserver mDarkModeObserver;
     private Runnable mCallback;
-    private Uri mUri = Settings.Secure.getUriFor(Settings.Secure.UI_NIGHT_MODE);
+
+    public DarkModeObserverTest(Uri uri) {
+        mUri = uri;
+    }
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = spy(RuntimeEnvironment.application);
+        mContext = spy(ApplicationProvider.getApplicationContext());
         mDarkModeObserver = new DarkModeObserver(mContext);
         mContentObserver = mDarkModeObserver.getContentObserver();
         mCallback = mock(Runnable.class);
@@ -51,16 +74,21 @@ public class DarkModeObserverTest {
     @Test
     public void callbackTest_subscribedCallbackCalled() {
         mDarkModeObserver.subscribe(mCallback);
+
         mContentObserver.onChange(false, mUri);
+
+        // 1x from subscribe. 1x from onChange
         Mockito.verify(mCallback, times(2)).run();
     }
 
     @Test
-    public void callbackTest_unsubscribedCallNotbackCalled() {
+    public void callbackTest_unsubscribedAfterSubscribe_noMoreCallbackCalled() {
         mDarkModeObserver.subscribe(mCallback);
-        mContentObserver.onChange(false, mUri);
         mDarkModeObserver.unsubscribe();
+
         mContentObserver.onChange(false, mUri);
-        Mockito.verify(mCallback, times(2)).run();
+
+        // 1x from subscribe.
+        Mockito.verify(mCallback).run();
     }
 }

tests/robotests/src/com/android/settings/display/darkmode/DarkModeScheduleSelectorControllerTest.java

@@ -28,13 +28,18 @@ import static org.mockito.Mockito.when;
 
 import android.app.UiModeManager;
 import android.content.Context;
+import android.content.pm.ActivityInfo;
+import android.content.res.Configuration;
+import android.content.res.Resources;
 import android.location.LocationManager;
 import android.os.PowerManager;
 
 import androidx.preference.DropDownPreference;
 import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
+import com.android.settings.testutils.BedtimeSettingsUtils;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -42,7 +47,6 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 @RunWith(RobolectricTestRunner.class)
 public class DarkModeScheduleSelectorControllerTest {
@@ -59,22 +63,39 @@ public class DarkModeScheduleSelectorControllerTest {
     private LocationManager mLocationManager;
     @Mock
     private PowerManager mPM;
+    @Mock
+    private Resources mResources;
+
+    private ActivityInfo mBedtimeActivityInfo = new ActivityInfo();
+    private Configuration mConfigNightYes = new Configuration();
+    private Configuration mConfigNightNo = new Configuration();
+    private BedtimeSettingsUtils mBedtimeSettingsUtils;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = spy(RuntimeEnvironment.application);
+        mContext = spy(ApplicationProvider.getApplicationContext());
+        mBedtimeSettingsUtils = new BedtimeSettingsUtils(mContext);
+        mConfigNightYes.uiMode = Configuration.UI_MODE_NIGHT_YES;
+        mConfigNightNo.uiMode = Configuration.UI_MODE_NIGHT_NO;
+
         when(mContext.getSystemService(UiModeManager.class)).thenReturn(mUiService);
         when(mContext.getSystemService(PowerManager.class)).thenReturn(mPM);
         when(mContext.getSystemService(LocationManager.class)).thenReturn(mLocationManager);
+        when(mContext.getResources()).thenReturn(mResources);
+        when(mResources.getConfiguration()).thenReturn(mConfigNightNo);
+        when(mResources.getString(com.android.internal.R.string.config_defaultWellbeingPackage))
+                .thenReturn("wellbeing");
+
         when(mContext.getString(R.string.dark_ui_auto_mode_never)).thenReturn("never");
         when(mContext.getString(R.string.dark_ui_auto_mode_auto)).thenReturn("auto");
         when(mContext.getString(R.string.dark_ui_auto_mode_custom)).thenReturn("custom");
+        when(mContext.getString(R.string.dark_ui_auto_mode_custom_bedtime))
+                .thenReturn("custom_bedtime");
+        when(mResources.getStringArray(R.array.dark_ui_scheduler_with_bedtime_preference_titles))
+                .thenReturn(new String[]{"never", "auto", "custom", "custom_bedtime"});
         mPreference = spy(new DropDownPreference(mContext));
-        mPreference.setEntryValues(new CharSequence[]{
-                mContext.getString(R.string.dark_ui_auto_mode_never),
-                mContext.getString(R.string.dark_ui_auto_mode_auto)
-        });
+        mPreference.setEntryValues(new CharSequence[]{"never", "auto", "custom"});
         doNothing().when(mPreference).setValueIndex(anyInt());
         when(mLocationManager.isLocationEnabled()).thenReturn(true);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
@@ -105,16 +126,74 @@ public class DarkModeScheduleSelectorControllerTest {
     @Test
     public void nightMode_updateStateNone_dropDownValueChangedToAuto() {
         when(mUiService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_AUTO);
+
         mController.displayPreference(mScreen);
         mController.updateState(mScreen);
+
         verify(mPreference).setValueIndex(1);
     }
 
     @Test
-    public void batterySaver_dropDown_disabledSelector() {
-        when(mPM.isPowerSaveMode()).thenReturn(true);
+    public void nightMode_updateStateNone_dropDownValueChangedToCustom() {
+        when(mUiService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+
         mController.displayPreference(mScreen);
         mController.updateState(mScreen);
+
+        verify(mPreference).setValueIndex(2);
+    }
+
+    @Test
+    public void nightMode_updateStateNone_bedtimeEnabled_dropDownValueChangedToCustomBedtime() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                true /* enabled */);
+        when(mUiService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mUiService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+        mBedtimeActivityInfo.enabled = true;
+
+        mController.displayPreference(mScreen);
+        mController.updateState(mScreen);
+
+        verify(mPreference).setValueIndex(3);
+    }
+
+    @Test
+    public void nightMode_updateStateNone_bedtimeDisabled_dropDownValueChangedToCustom() {
+        mBedtimeSettingsUtils.installBedtimeSettings("wellbeing" /* wellbeingPackage */,
+                false /* enabled */);
+        mController = new DarkModeScheduleSelectorController(mContext, mPreferenceKey);
+        when(mUiService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mUiService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+
+        mController.displayPreference(mScreen);
+        mController.updateState(mScreen);
+
+        verify(mPreference).setValueIndex(2);
+    }
+
+    @Test
+    public void nightMode_updateStateNone_bedtimeNotInstalled_dropDownValueChangedToCustom() {
+        mController = new DarkModeScheduleSelectorController(mContext, mPreferenceKey);
+        when(mUiService.getNightMode()).thenReturn(UiModeManager.MODE_NIGHT_CUSTOM);
+        when(mUiService.getNightModeCustomType())
+                .thenReturn(UiModeManager.MODE_NIGHT_CUSTOM_TYPE_BEDTIME);
+
+        mController.displayPreference(mScreen);
+        mController.updateState(mScreen);
+
+        verify(mPreference).setValueIndex(2);
+    }
+
+
+    @Test
+    public void batterySaver_dropDown_disabledSelector() {
+        when(mPM.isPowerSaveMode()).thenReturn(true);
+
+        mController.displayPreference(mScreen);
+        mController.updateState(mScreen);
+
         verify(mPreference).setEnabled(eq(false));
     }
 }

tests/robotests/src/com/android/settings/testutils/BedtimeSettingsUtils.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2022 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.settings.testutils;
+
+import static com.android.settings.display.darkmode.BedtimeSettings.ACTION_BEDTIME_SETTINGS;
+
+import static org.robolectric.Shadows.shadowOf;
+
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ActivityInfo;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.ResolveInfo;
+
+/** A helper class for installing bedtime settings activity. */
+public final class BedtimeSettingsUtils {
+    private Context mContext;
+
+    public BedtimeSettingsUtils(Context context) {
+        mContext = context;
+    }
+
+    public void installBedtimeSettings(String wellbeingPackage, boolean enabled) {
+        Intent bedtimeSettingsIntent = new Intent(ACTION_BEDTIME_SETTINGS)
+                .setPackage(wellbeingPackage);
+        ResolveInfo bedtimeResolveInfo = new ResolveInfo();
+        bedtimeResolveInfo.activityInfo = new ActivityInfo();
+        bedtimeResolveInfo.activityInfo.name = "BedtimeSettings";
+        bedtimeResolveInfo.activityInfo.packageName = "wellbeing";
+        bedtimeResolveInfo.activityInfo.enabled = enabled;
+        bedtimeResolveInfo.activityInfo.applicationInfo = new ApplicationInfo();
+        bedtimeResolveInfo.activityInfo.applicationInfo.enabled = true;
+        shadowOf(mContext.getPackageManager()).addResolveInfoForIntent(
+                bedtimeSettingsIntent, bedtimeResolveInfo);
+    }
+}