refactor(brightness suw): decouple auto brightness initialization logic for setup flow

The AccessibilitySettingsForSetupWizard class should not be responsible for managing the initial
logic of controllers, such as determining which controllers need to be dynamically added or updated
using setInSetupWizard APIs. This logic should be handled directly by the controllers themselves.

Bug: 311093618
Flag: EXEMPT bugfix
Test: atest AccessibilitySettingsForSetupWizardTest
            AutoBrightnessPreferenceControllerForSetupWizardTest
            AutoBrightnessPreferenceControllerTest
            BrightnessLevelPreferenceControllerForSetupWizardTest
            BrightnessLevelPreferenceControllerTest

Change-Id: I6065a10e72d002981c0f514543e6933d79c2aa1b
diff --git a/res/xml/accessibility_settings_for_setup_wizard.xml b/res/xml/accessibility_settings_for_setup_wizard.xml
index a69dc26..25b8a5a 100644
--- a/res/xml/accessibility_settings_for_setup_wizard.xml
+++ b/res/xml/accessibility_settings_for_setup_wizard.xml
@@ -43,7 +43,8 @@
         android:title="@string/auto_brightness_title"
         android:fragment="com.android.settings.accessibility.AutoBrightnessPreferenceFragmentForSetupWizard"
         settings:useAdminDisabledSummary="true"
-        settings:userRestriction="no_config_brightness"/>
+        settings:userRestriction="no_config_brightness"
+        settings:controller="com.android.settings.display.AutoBrightnessPreferenceControllerForSetupWizard"/>
 
     <Preference
         android:fragment="com.android.settings.accessibility.ToggleScreenMagnificationPreferenceFragmentForSetupWizard"
diff --git a/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizard.java b/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizard.java
index 447e34e..11cbd16 100644
--- a/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizard.java
+++ b/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizard.java
@@ -41,8 +41,7 @@
 
 import com.android.settings.R;
 import com.android.settings.dashboard.DashboardFragment;
-import com.android.settings.display.AutoBrightnessPreferenceController;
-import com.android.settings.display.BrightnessLevelPreferenceController;
+import com.android.settings.display.BrightnessLevelPreferenceControllerForSetupWizard;
 import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.core.AbstractPreferenceController;
 
@@ -168,16 +167,10 @@
 
     @Override
     protected List<AbstractPreferenceController> createPreferenceControllers(Context context) {
+        // Requires lifecycle, so added programmatically (normally via resId).
         final List<AbstractPreferenceController> controllers = new ArrayList<>();
-        BrightnessLevelPreferenceController brightnessLevelPreferenceController =
-                new BrightnessLevelPreferenceController(context, getSettingsLifecycle());
-        brightnessLevelPreferenceController.setInSetupWizard(true);
-        controllers.add(brightnessLevelPreferenceController);
-        String autoBrightnessKey = context.getString(R.string.preference_key_auto_brightness);
-        AutoBrightnessPreferenceController autoBrightnessPreferenceController =
-                new AutoBrightnessPreferenceController(context, autoBrightnessKey);
-        autoBrightnessPreferenceController.setInSetupWizard(true);
-        controllers.add(autoBrightnessPreferenceController);
+        controllers.add(new BrightnessLevelPreferenceControllerForSetupWizard(
+                context, getSettingsLifecycle()));
         return controllers;
     }
 
diff --git a/src/com/android/settings/display/AutoBrightnessPreferenceController.java b/src/com/android/settings/display/AutoBrightnessPreferenceController.java
index 3014f62..0a80d69 100644
--- a/src/com/android/settings/display/AutoBrightnessPreferenceController.java
+++ b/src/com/android/settings/display/AutoBrightnessPreferenceController.java
@@ -22,28 +22,25 @@
 import android.os.UserManager;
 import android.provider.Settings;
 
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
 
 import com.android.settings.R;
-import com.android.settings.accessibility.Flags;
 import com.android.settings.core.TogglePreferenceController;
 import com.android.settingslib.PrimarySwitchPreference;
 
+/**
+ * The top-level preference controller that updates the adaptive brightness.
+ */
 public class AutoBrightnessPreferenceController extends TogglePreferenceController {
 
     private final String SYSTEM_KEY = SCREEN_BRIGHTNESS_MODE;
     private final int DEFAULT_VALUE = SCREEN_BRIGHTNESS_MODE_MANUAL;
 
-    private boolean mInSetupWizard;
-
-    public AutoBrightnessPreferenceController(Context context, String key) {
+    public AutoBrightnessPreferenceController(@NonNull Context context, @NonNull String key) {
         super(context, key);
     }
 
-    public void setInSetupWizard(boolean inSetupWizard) {
-        mInSetupWizard = inSetupWizard;
-    }
-
     @Override
     public boolean isChecked() {
         return Settings.System.getInt(mContext.getContentResolver(),
@@ -60,14 +57,10 @@
     @Override
     @AvailabilityStatus
     public int getAvailabilityStatus() {
-        if (!mContext.getResources().getBoolean(
-                com.android.internal.R.bool.config_automatic_brightness_available)) {
-            return UNSUPPORTED_ON_DEVICE;
-        }
-        if (mInSetupWizard && !Flags.addBrightnessSettingsInSuw()) {
-            return CONDITIONALLY_UNAVAILABLE;
-        }
-        return AVAILABLE_UNSEARCHABLE;
+        return mContext.getResources().getBoolean(
+                com.android.internal.R.bool.config_automatic_brightness_available)
+                ? AVAILABLE_UNSEARCHABLE
+                : UNSUPPORTED_ON_DEVICE;
     }
 
     @Override
diff --git a/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizard.java b/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizard.java
new file mode 100644
index 0000000..7cdf3e0
--- /dev/null
+++ b/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizard.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 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;
+
+import android.content.Context;
+
+import androidx.annotation.NonNull;
+
+import com.android.settings.accessibility.Flags;
+
+/**
+ * The top-level preference controller that updates the adaptive brightness in the SetupWizard.
+ */
+public class AutoBrightnessPreferenceControllerForSetupWizard
+        extends AutoBrightnessPreferenceController {
+
+    public AutoBrightnessPreferenceControllerForSetupWizard(@NonNull Context context,
+            @NonNull String key) {
+        super(context, key);
+    }
+
+    @Override
+    @AvailabilityStatus
+    public int getAvailabilityStatus() {
+        if (!Flags.addBrightnessSettingsInSuw()) {
+            return CONDITIONALLY_UNAVAILABLE;
+        }
+        return super.getAvailabilityStatus();
+    }
+}
diff --git a/src/com/android/settings/display/BrightnessLevelPreferenceController.java b/src/com/android/settings/display/BrightnessLevelPreferenceController.java
index a32c965..96043db 100644
--- a/src/com/android/settings/display/BrightnessLevelPreferenceController.java
+++ b/src/com/android/settings/display/BrightnessLevelPreferenceController.java
@@ -20,7 +20,6 @@
 import static com.android.settingslib.display.BrightnessUtils.GAMMA_SPACE_MIN;
 import static com.android.settingslib.display.BrightnessUtils.convertLinearToGammaFloat;
 
-import android.annotation.Nullable;
 import android.app.ActivityOptions;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -37,12 +36,13 @@
 import android.provider.Settings.System;
 import android.text.TextUtils;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
 import com.android.settings.Utils;
-import com.android.settings.accessibility.Flags;
 import com.android.settings.core.BasePreferenceController;
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.core.SettingsBaseActivity;
@@ -54,18 +54,17 @@
 
 import java.text.NumberFormat;
 
+/**
+ * The top-level preference controller that updates the adaptive brightness level.
+ */
 public class BrightnessLevelPreferenceController extends BasePreferenceController implements
         PreferenceControllerMixin, LifecycleObserver, OnStart, OnStop {
-
-    private static final String TAG = "BrightnessPrefCtrl";
-
     private static final Uri BRIGHTNESS_ADJ_URI;
     private final ContentResolver mContentResolver;
     private final Handler mHandler = new Handler(Looper.getMainLooper());
     private final DisplayManager mDisplayManager;
     @Nullable
     private Preference mPreference;
-    private boolean mInSetupWizard;
 
     static {
         BRIGHTNESS_ADJ_URI = System.getUriFor(System.SCREEN_AUTO_BRIGHTNESS_ADJ);
@@ -94,11 +93,13 @@
         }
     };
 
-    public BrightnessLevelPreferenceController(Context context, Lifecycle lifecycle) {
+    public BrightnessLevelPreferenceController(@NonNull Context context,
+            @Nullable Lifecycle lifecycle) {
         this(context, context.getString(R.string.preference_key_brightness_level), lifecycle);
     }
 
-    private BrightnessLevelPreferenceController(Context context, String key, Lifecycle lifecycle) {
+    private BrightnessLevelPreferenceController(@NonNull Context context, @NonNull String key,
+            @Nullable Lifecycle lifecycle) {
         super(context, key);
         mDisplayManager = context.getSystemService(DisplayManager.class);
 
@@ -108,15 +109,8 @@
         mContentResolver = mContext.getContentResolver();
     }
 
-    public void setInSetupWizard(boolean inSetupWizard) {
-        mInSetupWizard = inSetupWizard;
-    }
-
     @Override
     public int getAvailabilityStatus() {
-        if (mInSetupWizard && !Flags.addBrightnessSettingsInSuw()) {
-            return CONDITIONALLY_UNAVAILABLE;
-        }
         return AVAILABLE;
     }
 
diff --git a/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizard.java b/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizard.java
new file mode 100644
index 0000000..6e002ae
--- /dev/null
+++ b/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizard.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 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;
+
+import android.content.Context;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.settings.accessibility.Flags;
+import com.android.settingslib.core.lifecycle.Lifecycle;
+
+/**
+ * The top-level preference controller that updates the adaptive brightness level in the
+ * SetupWizard.
+ */
+public class BrightnessLevelPreferenceControllerForSetupWizard extends
+        BrightnessLevelPreferenceController {
+
+    public BrightnessLevelPreferenceControllerForSetupWizard(@NonNull Context context,
+            @Nullable Lifecycle lifecycle) {
+        super(context, lifecycle);
+    }
+
+    @Override
+    @AvailabilityStatus
+    public int getAvailabilityStatus() {
+        if (!Flags.addBrightnessSettingsInSuw()) {
+            return CONDITIONALLY_UNAVAILABLE;
+        }
+        return super.getAvailabilityStatus();
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizardTest.java b/tests/robotests/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizardTest.java
index 2ae5984..58578e7 100644
--- a/tests/robotests/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizardTest.java
+++ b/tests/robotests/src/com/android/settings/accessibility/AccessibilitySettingsForSetupWizardTest.java
@@ -46,8 +46,7 @@
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
-import com.android.settings.display.AutoBrightnessPreferenceController;
-import com.android.settings.display.BrightnessLevelPreferenceController;
+import com.android.settings.display.BrightnessLevelPreferenceControllerForSetupWizard;
 import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.core.AbstractPreferenceController;
 
@@ -146,16 +145,15 @@
     }
 
     @Test
-    public void createPreferenceControllers_brightnessPreferencesControllersAreCreated() {
+    public void createPreferenceControllers_lifeCycleDependencyControllerIsCreated() {
         mFragment.onAttach(mContext);
 
         List<AbstractPreferenceController> controllers =
                 mFragment.createPreferenceControllers(mContext);
 
         assertTrue(controllers.stream().anyMatch(
-                controller -> controller instanceof BrightnessLevelPreferenceController));
-        assertTrue(controllers.stream().anyMatch(
-                controller -> controller instanceof AutoBrightnessPreferenceController));
+                controller ->
+                        controller instanceof BrightnessLevelPreferenceControllerForSetupWizard));
     }
 
     private void addEnabledServiceInfo(ComponentName componentName, boolean isAccessibilityTool) {
diff --git a/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizardTest.java b/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizardTest.java
new file mode 100644
index 0000000..1113d68
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerForSetupWizardTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2024 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;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE_UNSEARCHABLE;
+import static com.android.settings.core.BasePreferenceController.CONDITIONALLY_UNAVAILABLE;
+import static com.android.settings.core.BasePreferenceController.UNSUPPORTED_ON_DEVICE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import com.android.settings.accessibility.Flags;
+import com.android.settings.testutils.shadow.SettingsShadowResources;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+
+/**
+ * Tests for {@link AutoBrightnessPreferenceControllerForSetupWizard}.
+ */
+@RunWith(RobolectricTestRunner.class)
+@Config(shadows = {SettingsShadowResources.class})
+public class AutoBrightnessPreferenceControllerForSetupWizardTest {
+
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private static final String PREFERENCE_KEY = "auto_brightness";
+
+    private Context mContext;
+    private AutoBrightnessPreferenceControllerForSetupWizard mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mController =
+                new AutoBrightnessPreferenceControllerForSetupWizard(mContext, PREFERENCE_KEY);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
+    public void getAvailabilityStatus_configTrueAndFlagOn_shouldReturnAvailableUnsearchable() {
+        SettingsShadowResources.overrideResource(
+                com.android.internal.R.bool.config_automatic_brightness_available, true);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE_UNSEARCHABLE);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
+    public void getAvailabilityStatus_configFalseSetAndFlagOn_shouldReturnUnsupportedOnDevice() {
+        SettingsShadowResources.overrideResource(
+                com.android.internal.R.bool.config_automatic_brightness_available, false);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(UNSUPPORTED_ON_DEVICE);
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
+    public void getAvailabilityStatus_flagOff_shouldReturnConditionallyUnavailable() {
+        SettingsShadowResources.overrideResource(
+                com.android.internal.R.bool.config_automatic_brightness_available, true);
+
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
+    }
+}
+
diff --git a/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerTest.java
index 13cd86d..0122044 100644
--- a/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/AutoBrightnessPreferenceControllerTest.java
@@ -21,24 +21,18 @@
 import static android.provider.Settings.System.SCREEN_BRIGHTNESS_MODE_MANUAL;
 
 import static com.android.settings.core.BasePreferenceController.AVAILABLE_UNSEARCHABLE;
-import static com.android.settings.core.BasePreferenceController.CONDITIONALLY_UNAVAILABLE;
 import static com.android.settings.core.BasePreferenceController.UNSUPPORTED_ON_DEVICE;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import android.content.ContentResolver;
 import android.content.Context;
-import android.platform.test.annotations.DisableFlags;
-import android.platform.test.annotations.EnableFlags;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
 
 import com.android.settings.R;
-import com.android.settings.accessibility.Flags;
 import com.android.settings.testutils.shadow.SettingsShadowResources;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
@@ -46,13 +40,13 @@
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
+/**
+ * Tests for {@link AutoBrightnessPreferenceController}.
+ */
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = {SettingsShadowResources.class})
 public class AutoBrightnessPreferenceControllerTest {
 
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
-
     private static final String PREFERENCE_KEY = "auto_brightness";
 
     private Context mContext;
@@ -130,7 +124,7 @@
     }
 
     @Test
-    public void getAvailabilityStatusNotInSUW_configTrueSet_shouldReturnAvailableUnsearchable() {
+    public void getAvailabilityStatus_configTrueSet_shouldReturnAvailableUnsearchable() {
         SettingsShadowResources.overrideResource(
                 com.android.internal.R.bool.config_automatic_brightness_available, true);
 
@@ -138,27 +132,6 @@
     }
 
     @Test
-    @EnableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
-    public void getAvailabilityStatusInSUW_configTrueAndFlagOn_shouldReturnAvailableUnsearchable() {
-        SettingsShadowResources.overrideResource(
-                com.android.internal.R.bool.config_automatic_brightness_available, true);
-        mController.setInSetupWizard(true);
-
-        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE_UNSEARCHABLE);
-    }
-
-    @Test
-    @DisableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
-    public void
-            getAvailabilityStatusInSUW_configTrueAndFlagOff_shouldReturnConditionallyUnavailable() {
-        SettingsShadowResources.overrideResource(
-                com.android.internal.R.bool.config_automatic_brightness_available, true);
-        mController.setInSetupWizard(true);
-
-        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
-    }
-
-    @Test
     public void getAvailabilityStatus_configFalseSet_shouldReturnUnsupportedOnDevice() {
         SettingsShadowResources.overrideResource(
                 com.android.internal.R.bool.config_automatic_brightness_available, false);
diff --git a/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizardTest.java b/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizardTest.java
new file mode 100644
index 0000000..25ddd1e
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerForSetupWizardTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2024 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;
+
+import static com.android.settings.core.BasePreferenceController.AVAILABLE;
+import static com.android.settings.core.BasePreferenceController.CONDITIONALLY_UNAVAILABLE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import com.android.settings.accessibility.Flags;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+/**
+ * Tests for {@link BrightnessLevelPreferenceControllerForSetupWizard}.
+ */
+@RunWith(RobolectricTestRunner.class)
+public class BrightnessLevelPreferenceControllerForSetupWizardTest {
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private Context mContext;
+    private BrightnessLevelPreferenceControllerForSetupWizard mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mController = new BrightnessLevelPreferenceControllerForSetupWizard(mContext,
+                /* lifecycle= */ null);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
+    public void getAvailabilityStatus_flagOn_shouldReturnAvailable() {
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(AVAILABLE);
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
+    public void getAvailabilityStatus_flagOff_shouldReturnConditionallyUnavailable() {
+        assertThat(mController.getAvailabilityStatus()).isEqualTo(CONDITIONALLY_UNAVAILABLE);
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerTest.java
index 009ca95..eb83040 100644
--- a/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/BrightnessLevelPreferenceControllerTest.java
@@ -33,9 +33,6 @@
 import android.content.Intent;
 import android.hardware.display.BrightnessInfo;
 import android.os.PowerManager;
-import android.platform.test.annotations.DisableFlags;
-import android.platform.test.annotations.EnableFlags;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings.System;
 import android.view.Display;
 
@@ -43,12 +40,10 @@
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settings.accessibility.Flags;
 import com.android.settings.core.SettingsBaseActivity;
 import com.android.settingslib.transition.SettingsTransitionHelper;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -61,12 +56,11 @@
 import org.robolectric.shadows.ShadowApplication;
 import org.robolectric.shadows.ShadowContentResolver;
 
+/**
+ * Tests for {@link BrightnessLevelPreferenceController}.
+ */
 @RunWith(RobolectricTestRunner.class)
 public class BrightnessLevelPreferenceControllerTest {
-
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
-
     @Mock
     private PowerManager mPowerManager;
     @Mock
@@ -95,29 +89,15 @@
                 mPowerManager);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
         doReturn(mDisplay).when(mContext).getDisplay();
-        mController = spy(new BrightnessLevelPreferenceController(mContext, null));
+        mController = spy(new BrightnessLevelPreferenceController(mContext, /* lifecycle= */ null));
     }
 
     @Test
-    public void isAvailable_shouldAlwaysReturnTrueWhenNotInSetupWizard() {
+    public void isAvailable_shouldAlwaysReturnTrue() {
         assertThat(mController.isAvailable()).isTrue();
     }
 
     @Test
-    @EnableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
-    public void isAvailable_inSetupWizardAndFlagOn_shouldReturnTrue() {
-        mController.setInSetupWizard(true);
-        assertThat(mController.isAvailable()).isTrue();
-    }
-
-    @Test
-    @DisableFlags(Flags.FLAG_ADD_BRIGHTNESS_SETTINGS_IN_SUW)
-    public void isAvailable_inSetupWizardAndFlagOff_shouldReturnFalse() {
-        mController.setInSetupWizard(true);
-        assertThat(mController.isAvailable()).isFalse();
-    }
-
-    @Test
     public void onStart_shouldRegisterObserver() {
         BrightnessLevelPreferenceController controller =
                 new BrightnessLevelPreferenceController(mContext, null);