Persist dev option reboot dialog on rotation.

Use a ViewModel to store reboot dialog fragment state on rotations. Test: locally verified rotations on few dev options that use the reboot fragment Test: atest com.android.settings.development.RebootConfirmationDialogFragmentTest --iterations 25 (Passed) Bug: 356273849 Flag: EXEMPT bug fix Change-Id: I154170ea5fa91bcbec32af43853c5f99f8e8253d
2025-02-18 13:34:58 +00:00
parent ba7233ad09
commit cf33160dc5
9 changed files with 299 additions and 66 deletions
--- a/src/com/android/settings/development/RebootConfirmationDialogFragment.java
+++ b/src/com/android/settings/development/RebootConfirmationDialogFragment.java
@@ -21,10 +21,13 @@ import android.app.settings.SettingsEnums;
 import android.content.DialogInterface;
 import android.os.Bundle;

+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.appcompat.app.AlertDialog;
 import androidx.fragment.app.Fragment;
 import androidx.fragment.app.FragmentManager;
+import androidx.lifecycle.ViewModelProvider;

 import com.android.settings.R;
 import com.android.settings.core.instrumentation.InstrumentedDialogFragment;
@@ -33,11 +36,10 @@ import com.android.settings.core.instrumentation.InstrumentedDialogFragment;
 public class RebootConfirmationDialogFragment extends InstrumentedDialogFragment
        implements DialogInterface.OnClickListener, DialogInterface.OnDismissListener {

-    private static final String TAG = "FreeformPrefRebootDlg";
+    @VisibleForTesting
+    static final String TAG = "DevOptionRebootDlg";

-    private final int mMessageId;
-    private final int mCancelButtonId;
-    private final RebootConfirmationDialogHost mHost;
+    private RebootConfirmationDialogViewModel mViewModel;

    /** Show an instance of this dialog. */
    public static void show(Fragment fragment, int messageId, RebootConfirmationDialogHost host) {
@@ -50,55 +52,60 @@ public class RebootConfirmationDialogFragment extends InstrumentedDialogFragment
            int messageId,
            int cancelButtonId,
            RebootConfirmationDialogHost host) {
-        final FragmentManager manager = fragment.getActivity().getSupportFragmentManager();
+        final FragmentManager manager = fragment.requireActivity().getSupportFragmentManager();
        if (manager.findFragmentByTag(TAG) == null) {
            final RebootConfirmationDialogFragment dialog =
-                    new RebootConfirmationDialogFragment(messageId, cancelButtonId, host);
+                    new RebootConfirmationDialogFragment();
+            RebootConfirmationDialogViewModel mViewModel = new ViewModelProvider(
+                    fragment.requireActivity()).get(
+                    RebootConfirmationDialogViewModel.class);
+            mViewModel.setMessageId(messageId);
+            mViewModel.setCancelButtonId(cancelButtonId);
+            mViewModel.setHost(host);
            dialog.show(manager, TAG);
        }
    }

-    @VisibleForTesting
-    RebootConfirmationDialogFragment(
-            int messageId, int cancelButtonId, RebootConfirmationDialogHost host) {
-        mMessageId = messageId;
-        mCancelButtonId = cancelButtonId;
-        mHost = host;
-    }
-
    @Override
    public int getMetricsCategory() {
        return SettingsEnums.REBOOT_CONFIRMATION_DIALOG;
    }

    @Override
-    public Dialog onCreateDialog(Bundle savedInstances) {
-        return new AlertDialog.Builder(getActivity())
-                .setMessage(mMessageId)
+    public void onCreate(@Nullable Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+        mViewModel = new ViewModelProvider(requireActivity()).get(
+                RebootConfirmationDialogViewModel.class);
+    }
+
+    @Override
+    public Dialog onCreateDialog(@Nullable Bundle savedInstances) {
+        int messageId = mViewModel.getMessageId();
+        int cancelButtonId = mViewModel.getCancelButtonId();
+        return new AlertDialog.Builder(requireActivity())
+                .setMessage(messageId)
                .setPositiveButton(R.string.reboot_dialog_reboot_now, this)
-                .setNegativeButton(mCancelButtonId, this)
+                .setNegativeButton(cancelButtonId, this)
                .create();
    }

    @Override
    public void onClick(DialogInterface dialog, int which) {
+        RebootConfirmationDialogHost host = mViewModel.getHost();
+        if (host == null) return;
        if (which == DialogInterface.BUTTON_POSITIVE) {
-            mHost.onRebootConfirmed(getContext());
+            host.onRebootConfirmed(requireContext());
        } else {
-            mHost.onRebootCancelled();
+            host.onRebootCancelled();
        }
    }

    @Override
-    public void onDismiss(DialogInterface dialog) {
+    public void onDismiss(@NonNull DialogInterface dialog) {
        super.onDismiss(dialog);
-        mHost.onRebootDialogDismissed();
-    }
-
-    @Override
-    public void onPause() {
-        dismiss();
-
-        super.onPause();
+        RebootConfirmationDialogHost host = mViewModel.getHost();
+        if (host != null) {
+            host.onRebootDialogDismissed();
+        }
    }
 }
--- a/src/com/android/settings/development/RebootConfirmationDialogViewModel.java
+++ b/src/com/android/settings/development/RebootConfirmationDialogViewModel.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2025 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.development;
+
+import androidx.annotation.Nullable;
+import androidx.lifecycle.ViewModel;
+
+/**
+ * {@link ViewModel} for the reboot confirmation dialog.
+ *
+ * This class holds the data necessary to display a confirmation dialog for reboots.
+ */
+public class RebootConfirmationDialogViewModel extends ViewModel {
+    @Nullable
+    private RebootConfirmationDialogHost mHost = null;
+    private int mMessageId;
+    private int mCancelButtonId;
+
+    @Nullable
+    public RebootConfirmationDialogHost getHost() {
+        return mHost;
+    }
+
+    public void setHost(RebootConfirmationDialogHost mHost) {
+        this.mHost = mHost;
+    }
+
+    public int getMessageId() {
+        return mMessageId;
+    }
+
+    public void setMessageId(int mMessageId) {
+        this.mMessageId = mMessageId;
+    }
+
+    public int getCancelButtonId() {
+        return mCancelButtonId;
+    }
+
+    public void setCancelButtonId(int mCancelButtonId) {
+        this.mCancelButtonId = mCancelButtonId;
+    }
+}
--- a/tests/robotests/src/com/android/settings/development/RebootConfirmationDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/development/RebootConfirmationDialogFragmentTest.java
@@ -16,7 +16,22 @@

 package com.android.settings.development;

+import static android.os.Looper.getMainLooper;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.content.DialogInterface;
+import android.widget.TextView;
+
+import androidx.appcompat.app.AlertDialog;
+import androidx.fragment.app.Fragment;
 import androidx.fragment.app.FragmentActivity;
+import androidx.fragment.app.FragmentManager;
+import androidx.fragment.app.FragmentTransaction;
+import androidx.lifecycle.ViewModelProvider;

 import com.android.settings.R;

@@ -24,38 +39,124 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.shadows.androidx.fragment.FragmentController;
+import org.robolectric.shadows.ShadowDialog;

@RunWith(RobolectricTestRunner.class)
 public class RebootConfirmationDialogFragmentTest {

-    private RebootConfirmationDialogFragment mFragment;
-
    @Mock
-    RebootConfirmationDialogHost mRebootConfirmationDialogHost;
+    private RebootConfirmationDialogHost mHost;
+    private FragmentActivity mActivity;
+    private Fragment mFragment;
+    private FragmentManager mFragmentManager;
+    private RebootConfirmationDialogViewModel mViewModel;

    @Before
    public void setUp() {
        MockitoAnnotations.initMocks(this);
-        RebootConfirmationDialogFragment dialogFragment =
-                FragmentController.setupFragment(
-                        new RebootConfirmationDialogFragment(
-                                R.string.reboot_dialog_override_desktop_mode,
-                                R.string.reboot_dialog_reboot_later,
-                                mRebootConfirmationDialogHost),
-                        FragmentActivity.class,
-                        0 /* containerViewId= */, null /* bundle= */);
-
-        mFragment = Mockito.spy(dialogFragment);
+        mActivity = Robolectric.buildActivity(FragmentActivity.class).create().get();
+        mFragmentManager = mActivity.getSupportFragmentManager();
+        mFragment = new Fragment();
+        FragmentTransaction transaction = mFragmentManager.beginTransaction();
+        transaction.add(mFragment, "testFragment");
+        transaction.commit();
+        mFragmentManager.executePendingTransactions();
    }

    @Test
-    public void onPause_shouldDismissDialog() {
-        mFragment.onPause();
+    public void show_shouldCreateAndShowDialog() {
+        RebootConfirmationDialogFragment.show(mFragment,
+                R.string.reboot_dialog_override_desktop_mode, mHost);
+        shadowOf(getMainLooper()).idle();

-        Mockito.verify(mFragment).dismiss();
+        RebootConfirmationDialogFragment dialogFragment =
+                (RebootConfirmationDialogFragment) mFragmentManager.findFragmentByTag(
+                        RebootConfirmationDialogFragment.TAG);
+        assertThat(dialogFragment).isNotNull();
+        assertThat(dialogFragment.getShowsDialog()).isTrue();
+    }
+
+    @Test
+    public void show_shouldStoreViewModel() {
+        RebootConfirmationDialogFragment.show(mFragment,
+                R.string.reboot_dialog_override_desktop_mode, R.string.reboot_dialog_reboot_later,
+                mHost);
+        shadowOf(getMainLooper()).idle();
+
+        mViewModel = new ViewModelProvider(mActivity).get(RebootConfirmationDialogViewModel.class);
+        assertThat(mViewModel.getHost()).isEqualTo(mHost);
+        assertThat(mViewModel.getMessageId()).isEqualTo(
+                R.string.reboot_dialog_override_desktop_mode);
+        assertThat(mViewModel.getCancelButtonId()).isEqualTo(R.string.reboot_dialog_reboot_later);
+    }
+
+    @Test
+    public void onCreateDialog_shouldCreateAlertDialogFromViewModel() {
+        RebootConfirmationDialogFragment dialogFragment = new RebootConfirmationDialogFragment();
+        dialogFragment.show(mFragmentManager, RebootConfirmationDialogFragment.TAG);
+        shadowOf(getMainLooper()).idle();
+        // Set up ViewModel
+        mViewModel = new ViewModelProvider(mActivity).get(RebootConfirmationDialogViewModel.class);
+        mViewModel.setMessageId(R.string.reboot_dialog_override_desktop_mode);
+        mViewModel.setCancelButtonId(R.string.reboot_dialog_reboot_later);
+        mViewModel.setHost(mHost);
+
+        dialogFragment.onCreateDialog(null).show();
+        shadowOf(getMainLooper()).idle();
+
+        AlertDialog alertDialog = (AlertDialog) ShadowDialog.getLatestDialog();
+        TextView messageView = alertDialog.findViewById(android.R.id.message);
+        assertThat(messageView.getText().toString()).isEqualTo(
+                mActivity.getString(R.string.reboot_dialog_override_desktop_mode));
+        assertThat(alertDialog.getButton(
+                DialogInterface.BUTTON_POSITIVE).getText().toString()).isEqualTo(
+                mActivity.getString(R.string.reboot_dialog_reboot_now));
+        assertThat(alertDialog.getButton(
+                DialogInterface.BUTTON_NEGATIVE).getText().toString()).isEqualTo(
+                mActivity.getString(R.string.reboot_dialog_reboot_later));
+    }
+
+    @Test
+    public void onClick_positiveButton_shouldCallRebootConfirmed() {
+        RebootConfirmationDialogFragment dialogFragment = showDialog();
+        AlertDialog alertDialog = (AlertDialog) ShadowDialog.getLatestDialog();
+
+        dialogFragment.onClick(alertDialog, DialogInterface.BUTTON_POSITIVE);
+        verify(mHost).onRebootConfirmed(mActivity);
+    }
+
+    @Test
+    public void onClick_negativeButton_shouldCallRebootCancelled() {
+        RebootConfirmationDialogFragment dialogFragment = showDialog();
+        AlertDialog alertDialog = (AlertDialog) ShadowDialog.getLatestDialog();
+
+        dialogFragment.onClick(alertDialog, DialogInterface.BUTTON_NEGATIVE);
+        verify(mHost).onRebootCancelled();
+    }
+
+    @Test
+    public void onDismiss_shouldCallRebootDialogDismissed() {
+        RebootConfirmationDialogFragment dialogFragment = showDialog();
+
+        dialogFragment.onDismiss(null);
+        verify(mHost).onRebootDialogDismissed();
+    }
+
+    private RebootConfirmationDialogFragment showDialog() {
+        RebootConfirmationDialogFragment dialogFragment = new RebootConfirmationDialogFragment();
+        dialogFragment.show(mFragmentManager, RebootConfirmationDialogFragment.TAG);
+        shadowOf(getMainLooper()).idle();
+
+        mViewModel = new ViewModelProvider(mActivity).get(RebootConfirmationDialogViewModel.class);
+        mViewModel.setMessageId(R.string.reboot_dialog_override_desktop_mode);
+        mViewModel.setCancelButtonId(R.string.reboot_dialog_reboot_later);
+        mViewModel.setHost(mHost);
+
+        dialogFragment.onCreateDialog(null).show();
+        shadowOf(getMainLooper()).idle();
+        return dialogFragment;
    }
 }
--- a/tests/robotests/src/com/android/settings/development/RebootConfirmationDialogViewModelTest.java
+++ b/tests/robotests/src/com/android/settings/development/RebootConfirmationDialogViewModelTest.java
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2025 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.development;
+
+import static com.google.common.truth.Truth.assertThat;
+
+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 RebootConfirmationDialogViewModelTest {
+
+    @Mock
+    private RebootConfirmationDialogHost mHost;
+
+    private RebootConfirmationDialogViewModel mViewModel;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mViewModel = new RebootConfirmationDialogViewModel();
+    }
+
+    @Test
+    public void getHost_returnsSetHost() {
+        mViewModel.setHost(mHost);
+        assertThat(mViewModel.getHost()).isEqualTo(mHost);
+    }
+
+    @Test
+    public void getMessageId_returnsSetMessageId() {
+        mViewModel.setMessageId(123);
+        assertThat(mViewModel.getMessageId()).isEqualTo(123);
+    }
+
+    @Test
+    public void getCancelButtonId_returnsSetCancelButtonId() {
+        mViewModel.setCancelButtonId(456);
+        assertThat(mViewModel.getCancelButtonId()).isEqualTo(456);
+    }
+}
--- a/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopExperiencePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopExperiencePreferenceControllerTest.java
@@ -54,6 +54,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowSystemProperties;
@@ -74,8 +75,6 @@ public class DesktopExperiencePreferenceControllerTest {
    @Mock
    private DevelopmentSettingsDashboardFragment mFragment;
    @Mock
-    private FragmentActivity mActivity;
-    @Mock
    private FragmentManager mFragmentManager;
    @Mock
    private FragmentTransaction mTransaction;
@@ -88,10 +87,12 @@ public class DesktopExperiencePreferenceControllerTest {
    public void setup() {
        MockitoAnnotations.initMocks(this);

+        FragmentActivity activity = spy(Robolectric.buildActivity(
+                FragmentActivity.class).create().get());
        mContext = spy(ApplicationProvider.getApplicationContext());
        doReturn(mTransaction).when(mFragmentManager).beginTransaction();
-        doReturn(mFragmentManager).when(mActivity).getSupportFragmentManager();
-        doReturn(mActivity).when(mFragment).getActivity();
+        doReturn(mFragmentManager).when(activity).getSupportFragmentManager();
+        doReturn(activity).when(mFragment).requireActivity();

        mResources = spy(mContext.getResources());
        when(mContext.getResources()).thenReturn(mResources);
--- a/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopModePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopModePreferenceControllerTest.java
@@ -54,6 +54,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 import org.robolectric.shadows.ShadowSystemProperties;
@@ -76,8 +77,6 @@ public class DesktopModePreferenceControllerTest {
    @Mock
    private DevelopmentSettingsDashboardFragment mFragment;
    @Mock
-    private FragmentActivity mActivity;
-    @Mock
    private FragmentManager mFragmentManager;
    @Mock
    private FragmentTransaction mTransaction;
@@ -90,10 +89,13 @@ public class DesktopModePreferenceControllerTest {
    public void setup() {
        MockitoAnnotations.initMocks(this);

+        FragmentActivity activity = spy(Robolectric.buildActivity(
+                FragmentActivity.class).create().get());
+
        mContext = spy(ApplicationProvider.getApplicationContext());
        doReturn(mTransaction).when(mFragmentManager).beginTransaction();
-        doReturn(mFragmentManager).when(mActivity).getSupportFragmentManager();
-        doReturn(mActivity).when(mFragment).getActivity();
+        doReturn(mFragmentManager).when(activity).getSupportFragmentManager();
+        doReturn(activity).when(mFragment).requireActivity();

        mResources = spy(mContext.getResources());
        when(mContext.getResources()).thenReturn(mResources);
--- a/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopModeSecondaryDisplayPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/desktopexperience/DesktopModeSecondaryDisplayPreferenceControllerTest.java
@@ -54,6 +54,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
@@ -76,8 +77,6 @@ public class DesktopModeSecondaryDisplayPreferenceControllerTest {
    @Mock
    private DevelopmentSettingsDashboardFragment mFragment;
    @Mock
-    private FragmentActivity mActivity;
-    @Mock
    private FragmentManager mFragmentManager;
    @Mock
    private FragmentTransaction mTransaction;
@@ -89,12 +88,14 @@ public class DesktopModeSecondaryDisplayPreferenceControllerTest {
    @Before
    public void setup() throws Exception {
        MockitoAnnotations.initMocks(this);
+        FragmentActivity activity = spy(Robolectric.buildActivity(
+                FragmentActivity.class).create().get());
        mContext = spy(RuntimeEnvironment.application);
        mResources = spy(mContext.getResources());
        when(mContext.getResources()).thenReturn(mResources);
        doReturn(mTransaction).when(mFragmentManager).beginTransaction();
-        doReturn(mFragmentManager).when(mActivity).getSupportFragmentManager();
-        doReturn(mActivity).when(mFragment).getActivity();
+        doReturn(mFragmentManager).when(activity).getSupportFragmentManager();
+        doReturn(activity).when(mFragment).requireActivity();
        mController = new DesktopModeSecondaryDisplayPreferenceController(mContext, mFragment);
        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPreference);
        mController.displayPreference(mScreen);
--- a/tests/robotests/src/com/android/settings/development/desktopexperience/FreeformWindowsPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/desktopexperience/FreeformWindowsPreferenceControllerTest.java
@@ -25,6 +25,7 @@ import static com.google.common.truth.Truth.assertThat;

 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;

@@ -53,6 +54,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;

@@ -78,8 +80,6 @@ public class FreeformWindowsPreferenceControllerTest {
    @Mock
    private DevelopmentSettingsDashboardFragment mFragment;
    @Mock
-    private FragmentActivity mActivity;
-    @Mock
    private FragmentManager mFragmentManager;
    @Mock
    private FragmentTransaction mTransaction;
@@ -89,9 +89,11 @@ public class FreeformWindowsPreferenceControllerTest {
    @Before
    public void setup() {
        MockitoAnnotations.initMocks(this);
+        FragmentActivity activity = spy(Robolectric.buildActivity(
+                FragmentActivity.class).create().get());
        doReturn(mTransaction).when(mFragmentManager).beginTransaction();
-        doReturn(mFragmentManager).when(mActivity).getSupportFragmentManager();
-        doReturn(mActivity).when(mFragment).getActivity();
+        doReturn(mFragmentManager).when(activity).getSupportFragmentManager();
+        doReturn(activity).when(mFragment).requireActivity();
        doReturn(true).when(mResources).getBoolean(R.bool.config_isDesktopModeSupported);
        doReturn(true).when(mResources).getBoolean(R.bool.config_canInternalDisplayHostDesktops);
        mController = new FreeformWindowsPreferenceController(mContext, mFragment);
--- a/tests/robotests/src/com/android/settings/development/graphicsdriver/GraphicsDriverEnableAngleAsSystemDriverControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/graphicsdriver/GraphicsDriverEnableAngleAsSystemDriverControllerTest.java
@@ -25,6 +25,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;

@@ -45,6 +46,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
@@ -59,7 +61,6 @@ public class GraphicsDriverEnableAngleAsSystemDriverControllerTest {
    @Mock private PreferenceScreen mScreen;
    @Mock private SwitchPreference mPreference;
    @Mock private DevelopmentSettingsDashboardFragment mFragment;
-    @Mock private FragmentActivity mActivity;
    @Mock private FragmentManager mFragmentManager;
    @Mock private FragmentTransaction mTransaction;

@@ -70,11 +71,13 @@ public class GraphicsDriverEnableAngleAsSystemDriverControllerTest {
    @Before
    public void setUp() {
        MockitoAnnotations.initMocks(this);
+        FragmentActivity activity = spy(Robolectric.buildActivity(
+                FragmentActivity.class).create().get());
        mContext = RuntimeEnvironment.application;
        ShadowSystemProperties.override(PROPERTY_DEBUG_ANGLE_DEVELOPER_OPTION, "true");
        doReturn(mTransaction).when(mFragmentManager).beginTransaction();
-        doReturn(mFragmentManager).when(mActivity).getSupportFragmentManager();
-        doReturn(mActivity).when(mFragment).getActivity();
+        doReturn(mFragmentManager).when(activity).getSupportFragmentManager();
+        doReturn(activity).when(mFragment).requireActivity();
        mController = new GraphicsDriverEnableAngleAsSystemDriverController(mContext, mFragment);
        when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPreference);
        mController.displayPreference(mScreen);