diff --git a/src/com/android/settings/accessibility/LaunchAccessibilityActivityPreferenceFragment.java b/src/com/android/settings/accessibility/LaunchAccessibilityActivityPreferenceFragment.java
index ccc7d60f3c4..b2c3c4dbf8c 100644
--- a/src/com/android/settings/accessibility/LaunchAccessibilityActivityPreferenceFragment.java
+++ b/src/com/android/settings/accessibility/LaunchAccessibilityActivityPreferenceFragment.java
@@ -97,12 +97,6 @@ public class LaunchAccessibilityActivityPreferenceFragment extends
         mSettingsTitle = (mSettingsIntent == null) ? null : settingsTitle;
     }
 
-    @Override
-    public void onSettingsClicked(ShortcutPreference preference) {
-        super.onSettingsClicked(preference);
-        showDialog(DialogEnums.EDIT_SHORTCUT);
-    }
-
     @Override
     int getUserShortcutTypes() {
         return AccessibilityUtil.getUserShortcutTypesFromSettings(getPrefContext(),
diff --git a/src/com/android/settings/accessibility/ToggleAccessibilityServicePreferenceFragment.java b/src/com/android/settings/accessibility/ToggleAccessibilityServicePreferenceFragment.java
index f892c5088f0..30946d30afb 100644
--- a/src/com/android/settings/accessibility/ToggleAccessibilityServicePreferenceFragment.java
+++ b/src/com/android/settings/accessibility/ToggleAccessibilityServicePreferenceFragment.java
@@ -311,7 +311,11 @@ public class ToggleAccessibilityServicePreferenceFragment extends
 
     @Override
     public void onSettingsClicked(ShortcutPreference preference) {
-        super.onSettingsClicked(preference);
+        // Do not restore shortcut in shortcut chooser dialog when shortcutPreference is turned off.
+        mUserShortcutTypesCache = mShortcutPreference.isChecked()
+                ? getUserShortcutTypes(getPrefContext(), UserShortcutType.SOFTWARE)
+                : UserShortcutType.EMPTY;
+
         final boolean isServiceOnOrShortcutAdded = mShortcutPreference.isChecked()
                 || mToggleServiceDividerSwitchPreference.isChecked();
         showPopupDialog(isServiceOnOrShortcutAdded ? DialogEnums.EDIT_SHORTCUT
diff --git a/src/com/android/settings/accessibility/ToggleColorInversionPreferenceFragment.java b/src/com/android/settings/accessibility/ToggleColorInversionPreferenceFragment.java
index dfdcd9e9f52..a45cd49656e 100644
--- a/src/com/android/settings/accessibility/ToggleColorInversionPreferenceFragment.java
+++ b/src/com/android/settings/accessibility/ToggleColorInversionPreferenceFragment.java
@@ -122,12 +122,6 @@ public class ToggleColorInversionPreferenceFragment extends ToggleFeaturePrefere
         return R.string.help_url_color_inversion;
     }
 
-    @Override
-    public void onSettingsClicked(ShortcutPreference preference) {
-        super.onSettingsClicked(preference);
-        showDialog(DialogEnums.EDIT_SHORTCUT);
-    }
-
     @Override
     int getUserShortcutTypes() {
         return AccessibilityUtil.getUserShortcutTypesFromSettings(getPrefContext(),
diff --git a/src/com/android/settings/accessibility/ToggleDaltonizerPreferenceFragment.java b/src/com/android/settings/accessibility/ToggleDaltonizerPreferenceFragment.java
index 6c47dc3f15f..3fb9b38e5c2 100644
--- a/src/com/android/settings/accessibility/ToggleDaltonizerPreferenceFragment.java
+++ b/src/com/android/settings/accessibility/ToggleDaltonizerPreferenceFragment.java
@@ -194,12 +194,6 @@ public final class ToggleDaltonizerPreferenceFragment extends ToggleFeaturePrefe
         });
     }
 
-    @Override
-    public void onSettingsClicked(ShortcutPreference preference) {
-        super.onSettingsClicked(preference);
-        showDialog(DialogEnums.EDIT_SHORTCUT);
-    }
-
     @Override
     int getUserShortcutTypes() {
         return AccessibilityUtil.getUserShortcutTypesFromSettings(getPrefContext(),
diff --git a/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragment.java b/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragment.java
index 56f67335944..88f58be066c 100644
--- a/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragment.java
+++ b/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragment.java
@@ -57,6 +57,8 @@ import com.android.settings.widget.SwitchBar;
 import com.android.settingslib.accessibility.AccessibilityUtils;
 import com.android.settingslib.widget.FooterPreference;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
@@ -87,19 +89,24 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
     protected Uri mImageUri;
     private CharSequence mDescription;
     protected CharSequence mHtmlDescription;
-    // Used to restore the edit dialog status.
-    protected int mUserShortcutTypesCache = UserShortcutType.EMPTY;
+
     private static final String DRAWABLE_FOLDER = "drawable";
     protected static final String KEY_USE_SERVICE_PREFERENCE = "use_service";
     protected static final String KEY_GENERAL_CATEGORY = "general_categories";
     protected static final String KEY_INTRODUCTION_CATEGORY = "introduction_categories";
     private static final String KEY_SHORTCUT_PREFERENCE = "shortcut_preference";
-    private static final String EXTRA_SHORTCUT_TYPE = "shortcut_type";
+    @VisibleForTesting
+    static final String EXTRA_SHORTCUT_TYPE = "shortcut_type";
+
     private TouchExplorationStateChangeListener mTouchExplorationStateChangeListener;
-    private int mUserShortcutTypes = UserShortcutType.EMPTY;
+    private SettingsContentObserver mSettingsContentObserver;
+
     private CheckBox mSoftwareTypeCheckBox;
     private CheckBox mHardwareTypeCheckBox;
-    private SettingsContentObserver mSettingsContentObserver;
+
+    // Used to restore the edit dialog status.
+    protected int mUserShortcutTypesCache = UserShortcutType.EMPTY;
+    protected int mUserShortcutTypes = UserShortcutType.EMPTY;
 
     // For html description of accessibility service, must follow the rule, such as
     // <img src="R.drawable.fileName"/>, a11y settings will get the resources successfully.
@@ -121,6 +128,12 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
     @Override
     public void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
+        // Restore the user shortcut type.
+        if (savedInstanceState != null && savedInstanceState.containsKey(EXTRA_SHORTCUT_TYPE)) {
+            mUserShortcutTypesCache = savedInstanceState.getInt(EXTRA_SHORTCUT_TYPE,
+                    UserShortcutType.EMPTY);
+        }
+
         setupDefaultShortcutIfNecessary(getPrefContext());
         final int resId = getPreferenceScreenResId();
         if (resId <= 0) {
@@ -150,7 +163,7 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
         initAnimatedImagePreference();
         initToggleServiceDividerSwitchPreference();
         initGeneralCategory();
-        initShortcutPreference(savedInstanceState);
+        initShortcutPreference();
         initSettingsPreference();
         initHtmlTextPreference();
         initFooterPreference();
@@ -521,7 +534,8 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
         });
     }
 
-    private void initializeDialogCheckBox(Dialog dialog) {
+    @VisibleForTesting
+    void initializeDialogCheckBox(Dialog dialog) {
         final View dialogSoftwareView = dialog.findViewById(R.id.software_shortcut);
         mSoftwareTypeCheckBox = dialogSoftwareView.findViewById(R.id.checkbox);
         setDialogTextAreaClickListener(dialogSoftwareView, mSoftwareTypeCheckBox);
@@ -544,7 +558,8 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
         checkBox.setChecked((mUserShortcutTypesCache & type) == type);
     }
 
-    private void updateUserShortcutType(boolean saveChanges) {
+    @VisibleForTesting
+    void updateUserShortcutType(boolean saveChanges) {
         mUserShortcutTypesCache = UserShortcutType.EMPTY;
         if (mSoftwareTypeCheckBox.isChecked()) {
             mUserShortcutTypesCache |= UserShortcutType.SOFTWARE;
@@ -677,13 +692,7 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
         }
     }
 
-    protected void initShortcutPreference(Bundle savedInstanceState) {
-        // Restore the user shortcut type.
-        if (savedInstanceState != null && savedInstanceState.containsKey(EXTRA_SHORTCUT_TYPE)) {
-            mUserShortcutTypesCache = savedInstanceState.getInt(EXTRA_SHORTCUT_TYPE,
-                    UserShortcutType.EMPTY);
-        }
-
+    protected void initShortcutPreference() {
         // Initial the shortcut preference.
         mShortcutPreference = new ShortcutPreference(getPrefContext(), null);
         mShortcutPreference.setPersistent(false);
@@ -737,6 +746,7 @@ public abstract class ToggleFeaturePreferenceFragment extends SettingsPreference
         mUserShortcutTypesCache = mShortcutPreference.isChecked()
                 ? getUserShortcutTypes(getPrefContext(), UserShortcutType.SOFTWARE)
                 : UserShortcutType.EMPTY;
+        showDialog(DialogEnums.EDIT_SHORTCUT);
     }
 
     private void createFooterPreference(CharSequence title) {
diff --git a/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragment.java b/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragment.java
index 6a6df332223..108c1c0405e 100644
--- a/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragment.java
+++ b/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragment.java
@@ -62,7 +62,7 @@ public class ToggleScreenMagnificationPreferenceFragment extends
 
     private static final String EXTRA_SHORTCUT_TYPE = "shortcut_type";
     private TouchExplorationStateChangeListener mTouchExplorationStateChangeListener;
-    private int mUserShortcutType = UserShortcutType.EMPTY;
+
     private CheckBox mSoftwareTypeCheckBox;
     private CheckBox mHardwareTypeCheckBox;
     private CheckBox mTripleTapTypeCheckBox;
@@ -158,7 +158,8 @@ public class ToggleScreenMagnificationPreferenceFragment extends
         generalCategory.addPreference(mSettingsPreference);
     }
 
-    private void initializeDialogCheckBox(AlertDialog dialog) {
+    @VisibleForTesting
+    void initializeDialogCheckBox(AlertDialog dialog) {
         final View dialogSoftwareView = dialog.findViewById(R.id.software_shortcut);
         mSoftwareTypeCheckBox = dialogSoftwareView.findViewById(R.id.checkbox);
         setDialogTextAreaClickListener(dialogSoftwareView, mSoftwareTypeCheckBox);
@@ -193,7 +194,8 @@ public class ToggleScreenMagnificationPreferenceFragment extends
         checkBox.setChecked((mUserShortcutTypesCache & type) == type);
     }
 
-    private void updateUserShortcutType(boolean saveChanges) {
+    @VisibleForTesting
+    void updateUserShortcutType(boolean saveChanges) {
         mUserShortcutTypesCache = UserShortcutType.EMPTY;
         if (mSoftwareTypeCheckBox.isChecked()) {
             mUserShortcutTypesCache |= UserShortcutType.SOFTWARE;
@@ -210,7 +212,7 @@ public class ToggleScreenMagnificationPreferenceFragment extends
             if (isChanged) {
                 setUserShortcutType(getPrefContext(), mUserShortcutTypesCache);
             }
-            mUserShortcutType = mUserShortcutTypesCache;
+            mUserShortcutTypes = mUserShortcutTypesCache;
         }
     }
 
@@ -289,9 +291,9 @@ public class ToggleScreenMagnificationPreferenceFragment extends
     @Override
     protected void callOnAlertDialogCheckboxClicked(DialogInterface dialog, int which) {
         updateUserShortcutType(/* saveChanges= */ true);
-        optInAllMagnificationValuesToSettings(getPrefContext(), mUserShortcutType);
-        optOutAllMagnificationValuesFromSettings(getPrefContext(), ~mUserShortcutType);
-        mShortcutPreference.setChecked(mUserShortcutType != UserShortcutType.EMPTY);
+        optInAllMagnificationValuesToSettings(getPrefContext(), mUserShortcutTypes);
+        optOutAllMagnificationValuesFromSettings(getPrefContext(), ~mUserShortcutTypes);
+        mShortcutPreference.setChecked(mUserShortcutTypes != UserShortcutType.EMPTY);
         mShortcutPreference.setSummary(
                 getShortcutTypeSummary(getPrefContext()));
     }
@@ -361,17 +363,17 @@ public class ToggleScreenMagnificationPreferenceFragment extends
     @Override
     protected void updateShortcutPreferenceData() {
         // Get the user shortcut type from settings provider.
-        mUserShortcutType = getUserShortcutTypeFromSettings(getPrefContext());
-        if (mUserShortcutType != UserShortcutType.EMPTY) {
-            setUserShortcutType(getPrefContext(), mUserShortcutType);
+        mUserShortcutTypes = getUserShortcutTypeFromSettings(getPrefContext());
+        if (mUserShortcutTypes != UserShortcutType.EMPTY) {
+            setUserShortcutType(getPrefContext(), mUserShortcutTypes);
         } else {
             //  Get the user shortcut type from shared_prefs if cannot get from settings provider.
-            mUserShortcutType = getUserShortcutTypes(getPrefContext(), UserShortcutType.SOFTWARE);
+            mUserShortcutTypes = getUserShortcutTypes(getPrefContext(), UserShortcutType.SOFTWARE);
         }
     }
 
     @Override
-    protected void initShortcutPreference(Bundle savedInstanceState) {
+    protected void initShortcutPreference() {
         mShortcutPreference = new ShortcutPreference(getPrefContext(), null);
         mShortcutPreference.setPersistent(false);
         mShortcutPreference.setKey(getShortcutPreferenceKey());
diff --git a/tests/robotests/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragmentTest.java b/tests/robotests/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragmentTest.java
index 3df9c048637..f008b03a752 100644
--- a/tests/robotests/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/accessibility/ToggleFeaturePreferenceFragmentTest.java
@@ -16,44 +16,87 @@
 
 package com.android.settings.accessibility;
 
+import static com.android.settings.accessibility.ToggleFeaturePreferenceFragment.EXTRA_SHORTCUT_TYPE;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.DialogInterface;
 import android.os.Bundle;
+import android.provider.Settings;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 
 import androidx.annotation.XmlRes;
+import androidx.appcompat.app.AlertDialog;
 import androidx.fragment.app.FragmentActivity;
+import androidx.preference.PreferenceManager;
+import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
+import com.android.settings.accessibility.AccessibilityUtil.UserShortcutType;
 import com.android.settings.accessibility.ToggleFeaturePreferenceFragment.AccessibilityUserShortcutType;
+import com.android.settings.testutils.shadow.ShadowFragment;
 
+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;
+import org.robolectric.annotation.Config;
 import org.robolectric.shadows.androidx.fragment.FragmentController;
 
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+/** Tests for {@link ToggleFeaturePreferenceFragment} */
 @RunWith(RobolectricTestRunner.class)
 public class ToggleFeaturePreferenceFragmentTest {
 
-    private ToggleFeaturePreferenceFragmentTestable mFragment;
-
     private static final String TEST_SERVICE_KEY_1 = "abc:111";
     private static final String TEST_SERVICE_KEY_2 = "mno:222";
     private static final String TEST_SERVICE_KEY_3 = "xyz:333";
-
     private static final String TEST_SERVICE_NAME_1 = "abc";
     private static final int TEST_SERVICE_VALUE_1 = 111;
 
+    private static final String PLACEHOLDER_PACKAGE_NAME = "com.placeholder.example";
+    private static final String PLACEHOLDER_CLASS_NAME = PLACEHOLDER_PACKAGE_NAME + ".placeholder";
+    private static final ComponentName PLACEHOLDER_COMPONENT_NAME = new ComponentName(
+            PLACEHOLDER_PACKAGE_NAME, PLACEHOLDER_CLASS_NAME);
+
+    private static final String SOFTWARE_SHORTCUT_KEY =
+            Settings.Secure.ACCESSIBILITY_BUTTON_TARGETS;
+    private static final String HARDWARE_SHORTCUT_KEY =
+            Settings.Secure.ACCESSIBILITY_SHORTCUT_TARGET_SERVICE;
+
+    private TestToggleFeaturePreferenceFragment mFragment;
+    private Context mContext = ApplicationProvider.getApplicationContext();
+
+    @Mock
+    private PreferenceManager mPreferenceManager;
+
+    @Before
+    public void setUpTestFragment() {
+        MockitoAnnotations.initMocks(this);
+
+        mFragment = spy(new TestToggleFeaturePreferenceFragment());
+        when(mFragment.getPreferenceManager()).thenReturn(mPreferenceManager);
+        when(mFragment.getPreferenceManager().getContext()).thenReturn(mContext);
+        when(mFragment.getContext()).thenReturn(mContext);
+        doReturn(null).when(mFragment).getPreferenceScreen();
+    }
+
     @Test
     public void a11yUserShortcutType_setConcatString_shouldReturnTargetValue() {
         final AccessibilityUserShortcutType shortcut = new AccessibilityUserShortcutType(
@@ -109,7 +152,6 @@ public class ToggleFeaturePreferenceFragmentTest {
 
     @Test
     public void createFragment_shouldOnlyAddPreferencesOnce() {
-        mFragment = spy(new ToggleFeaturePreferenceFragmentTestable());
         FragmentController.setupFragment(mFragment, FragmentActivity.class,
                 /* containerViewId= */ 0, /* bundle= */null);
 
@@ -117,7 +159,75 @@ public class ToggleFeaturePreferenceFragmentTest {
         verify(mFragment).addPreferencesFromResource(R.xml.placeholder_prefs);
     }
 
-    public static class ToggleFeaturePreferenceFragmentTestable
+    @Test
+    public void updateShortcutPreferenceData_assignDefaultValueToVariable() {
+        mFragment.mComponentName = PLACEHOLDER_COMPONENT_NAME;
+
+        mFragment.updateShortcutPreferenceData();
+
+        // Compare to default UserShortcutType
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(UserShortcutType.SOFTWARE);
+    }
+
+    @Test
+    public void updateShortcutPreferenceData_hasValueInSettings_assignToVariable() {
+        mFragment.mComponentName = PLACEHOLDER_COMPONENT_NAME;
+
+        putStringIntoSettings(SOFTWARE_SHORTCUT_KEY, PLACEHOLDER_COMPONENT_NAME.flattenToString());
+        putStringIntoSettings(HARDWARE_SHORTCUT_KEY, PLACEHOLDER_COMPONENT_NAME.flattenToString());
+        mFragment.updateShortcutPreferenceData();
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(
+                UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE);
+    }
+
+    @Test
+    public void updateShortcutPreferenceData_hasValueInSharedPreference_assignToVariable() {
+        mFragment.mComponentName = PLACEHOLDER_COMPONENT_NAME;
+        final AccessibilityUserShortcutType hardwareShortcut = new AccessibilityUserShortcutType(
+                PLACEHOLDER_COMPONENT_NAME.flattenToString(), UserShortcutType.HARDWARE);
+
+        putUserShortcutTypeIntoSharedPreference(mContext, hardwareShortcut);
+        mFragment.updateShortcutPreferenceData();
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(UserShortcutType.HARDWARE);
+    }
+
+    @Test
+    @Config(shadows = ShadowFragment.class)
+    public void restoreValueFromSavedInstanceState_assignToVariable() {
+        mContext.setTheme(R.style.Theme_AppCompat);
+        final String dialogTitle = "title";
+        final AlertDialog dialog = AccessibilityEditDialogUtils.showEditShortcutDialog(
+                mContext, dialogTitle, this::callEmptyOnClicked);
+        final Bundle savedInstanceState = new Bundle();
+        mFragment.mComponentName = PLACEHOLDER_COMPONENT_NAME;
+
+        savedInstanceState.putInt(EXTRA_SHORTCUT_TYPE,
+                UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE);
+        mFragment.onCreate(savedInstanceState);
+        mFragment.initializeDialogCheckBox(dialog);
+        mFragment.updateUserShortcutType(true);
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(
+                UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE);
+
+    }
+
+    private void putStringIntoSettings(String key, String componentName) {
+        Settings.Secure.putString(mContext.getContentResolver(), key, componentName);
+    }
+
+    private void putUserShortcutTypeIntoSharedPreference(Context context,
+            AccessibilityUserShortcutType shortcut) {
+        Set<String> value = new HashSet<>(Collections.singletonList(shortcut.flattenToString()));
+
+        SharedPreferenceUtils.setUserShortcutType(context, value);
+    }
+
+    private void callEmptyOnClicked(DialogInterface dialog, int which) {}
+
+    public static class TestToggleFeaturePreferenceFragment
             extends ToggleFeaturePreferenceFragment {
 
         @Override
@@ -145,14 +255,25 @@ public class ToggleFeaturePreferenceFragmentTest {
             return mock(View.class);
         }
 
+        @Override
+        public void onViewCreated(View view, Bundle savedInstanceState) {
+            // do nothing
+        }
+
+        @Override
+        public void onDestroyView() {
+            // do nothing
+        }
+
         @Override
         public void addPreferencesFromResource(@XmlRes int preferencesResId) {
             // do nothing
         }
 
         @Override
-        public void onViewCreated(View view, Bundle savedInstanceState) {
-            // do nothing
+        protected void updateShortcutPreference() {
+            // UI related function, do nothing in tests
         }
+
     }
 }
diff --git a/tests/robotests/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragmentTest.java b/tests/robotests/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragmentTest.java
index 10a4d85b245..8bc3009d2fb 100644
--- a/tests/robotests/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/accessibility/ToggleScreenMagnificationPreferenceFragmentTest.java
@@ -19,40 +19,83 @@ package com.android.settings.accessibility;
 import static com.android.settings.accessibility.AccessibilityUtil.State.OFF;
 import static com.android.settings.accessibility.AccessibilityUtil.State.ON;
 import static com.android.settings.accessibility.AccessibilityUtil.UserShortcutType;
+import static com.android.settings.accessibility.ToggleFeaturePreferenceFragment.EXTRA_SHORTCUT_TYPE;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.DialogInterface;
+import android.os.Bundle;
 import android.provider.Settings;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.ViewGroup;
+
+import androidx.annotation.XmlRes;
+import androidx.appcompat.app.AlertDialog;
+import androidx.fragment.app.FragmentActivity;
+import androidx.preference.PreferenceManager;
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settings.R;
+import com.android.settings.accessibility.ToggleFeaturePreferenceFragment.AccessibilityUserShortcutType;
+import com.android.settings.testutils.shadow.ShadowFragment;
 
 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;
-import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 
 @RunWith(RobolectricTestRunner.class)
 public class ToggleScreenMagnificationPreferenceFragmentTest {
 
-    private static final String DUMMY_PACKAGE_NAME = "com.mock.example";
-    private static final String DUMMY_CLASS_NAME = DUMMY_PACKAGE_NAME + ".mock_a11y_service";
-    private static final ComponentName DUMMY_COMPONENT_NAME = new ComponentName(DUMMY_PACKAGE_NAME,
-            DUMMY_CLASS_NAME);
+    private static final String PLACEHOLDER_PACKAGE_NAME = "com.mock.example";
+    private static final String PLACEHOLDER_CLASS_NAME =
+            PLACEHOLDER_PACKAGE_NAME + ".mock_a11y_service";
+    private static final ComponentName PLACEHOLDER_COMPONENT_NAME = new ComponentName(
+            PLACEHOLDER_PACKAGE_NAME, PLACEHOLDER_CLASS_NAME);
+
     private static final String SOFTWARE_SHORTCUT_KEY =
             Settings.Secure.ACCESSIBILITY_BUTTON_TARGETS;
     private static final String HARDWARE_SHORTCUT_KEY =
             Settings.Secure.ACCESSIBILITY_SHORTCUT_TARGET_SERVICE;
     private static final String TRIPLETAP_SHORTCUT_KEY =
             Settings.Secure.ACCESSIBILITY_DISPLAY_MAGNIFICATION_ENABLED;
+
     private static final String MAGNIFICATION_CONTROLLER_NAME =
             "com.android.server.accessibility.MagnificationController";
 
-    private Context mContext;
+    private TestToggleScreenMagnificationPreferenceFragment mFragment;
+    private Context mContext = ApplicationProvider.getApplicationContext();
+
+    @Mock
+    private PreferenceManager mPreferenceManager;
+    @Mock
+    private FragmentActivity mActivity;
 
     @Before
-    public void setUp() {
-        mContext = RuntimeEnvironment.application;
+    public void setUpTestFragment() {
+        MockitoAnnotations.initMocks(this);
+
+        mFragment = spy(new TestToggleScreenMagnificationPreferenceFragment());
+        when(mFragment.getPreferenceManager()).thenReturn(mPreferenceManager);
+        when(mFragment.getPreferenceManager().getContext()).thenReturn(mContext);
+        when(mFragment.getContext()).thenReturn(mContext);
+        doReturn(null).when(mFragment).getPreferenceScreen();
+        doReturn(mActivity).when(mFragment).getActivity();
     }
 
     @Test
@@ -78,14 +121,13 @@ public class ToggleScreenMagnificationPreferenceFragmentTest {
 
     @Test
     public void optInAllValuesToSettings_existOtherValue_optInValue_haveMatchString() {
-        putStringIntoSettings(SOFTWARE_SHORTCUT_KEY, DUMMY_COMPONENT_NAME.flattenToString());
+        putStringIntoSettings(SOFTWARE_SHORTCUT_KEY, PLACEHOLDER_COMPONENT_NAME.flattenToString());
 
         ToggleScreenMagnificationPreferenceFragment.optInAllMagnificationValuesToSettings(mContext,
                 UserShortcutType.SOFTWARE);
 
         assertThat(getStringFromSettings(SOFTWARE_SHORTCUT_KEY)).isEqualTo(
-                DUMMY_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
-
+                PLACEHOLDER_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
     }
 
     @Test
@@ -107,27 +149,83 @@ public class ToggleScreenMagnificationPreferenceFragmentTest {
     @Test
     public void optOutValueFromSettings_existOtherValue_optOutValue_haveMatchString() {
         putStringIntoSettings(SOFTWARE_SHORTCUT_KEY,
-                DUMMY_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
+                PLACEHOLDER_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
         putStringIntoSettings(HARDWARE_SHORTCUT_KEY,
-                DUMMY_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
+                PLACEHOLDER_COMPONENT_NAME.flattenToString() + ":" + MAGNIFICATION_CONTROLLER_NAME);
         int shortcutTypes = UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE;
 
         ToggleScreenMagnificationPreferenceFragment.optOutAllMagnificationValuesFromSettings(
                 mContext, shortcutTypes);
 
         assertThat(getStringFromSettings(SOFTWARE_SHORTCUT_KEY)).isEqualTo(
-                DUMMY_COMPONENT_NAME.flattenToString());
+                PLACEHOLDER_COMPONENT_NAME.flattenToString());
         assertThat(getStringFromSettings(HARDWARE_SHORTCUT_KEY)).isEqualTo(
-                DUMMY_COMPONENT_NAME.flattenToString());
+                PLACEHOLDER_COMPONENT_NAME.flattenToString());
+    }
+
+    @Test
+    public void updateShortcutPreferenceData_assignDefaultValueToVariable() {
+        mFragment.updateShortcutPreferenceData();
+
+        // Compare to default UserShortcutType
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(UserShortcutType.SOFTWARE);
+    }
+
+    @Test
+    public void updateShortcutPreferenceData_hasValueInSettings_assignToVariable() {
+        putStringIntoSettings(SOFTWARE_SHORTCUT_KEY, MAGNIFICATION_CONTROLLER_NAME);
+        setMagnificationTripleTapEnabled(/* enabled= */ true);
+        mFragment.updateShortcutPreferenceData();
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(
+                UserShortcutType.SOFTWARE | UserShortcutType.TRIPLETAP);
+    }
+
+    @Test
+    public void updateShortcutPreferenceData_hasValueInSharedPreference_assignToVariable() {
+        final AccessibilityUserShortcutType tripleTapShortcut = new AccessibilityUserShortcutType(
+                MAGNIFICATION_CONTROLLER_NAME, UserShortcutType.TRIPLETAP);
+
+        putUserShortcutTypeIntoSharedPreference(mContext, tripleTapShortcut);
+        mFragment.updateShortcutPreferenceData();
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(UserShortcutType.TRIPLETAP);
+    }
+
+    @Test
+    @Config(shadows = ShadowFragment.class)
+    public void restoreValueFromSavedInstanceState_assignToVariable() {
+        mContext.setTheme(R.style.Theme_AppCompat);
+        final String dialogTitle = "title";
+        final AlertDialog dialog = AccessibilityEditDialogUtils.showMagnificationEditShortcutDialog(
+                mContext, dialogTitle, this::callEmptyOnClicked);
+        final Bundle savedInstanceState = new Bundle();
+
+        savedInstanceState.putInt(EXTRA_SHORTCUT_TYPE,
+                UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE);
+        mFragment.onCreate(savedInstanceState);
+        mFragment.initializeDialogCheckBox(dialog);
+        mFragment.updateUserShortcutType(true);
+
+        assertThat(mFragment.mUserShortcutTypes).isEqualTo(
+                UserShortcutType.SOFTWARE | UserShortcutType.HARDWARE);
+
     }
 
     private void putStringIntoSettings(String key, String componentName) {
         Settings.Secure.putString(mContext.getContentResolver(), key, componentName);
     }
 
+    private void putUserShortcutTypeIntoSharedPreference(Context context,
+            AccessibilityUserShortcutType shortcut) {
+        Set<String> value = new HashSet<>(Collections.singletonList(shortcut.flattenToString()));
+
+        SharedPreferenceUtils.setUserShortcutType(context, value);
+    }
+
     private void setMagnificationTripleTapEnabled(boolean enabled) {
-        Settings.Secure.putInt(mContext.getContentResolver(),
-                Settings.Secure.ACCESSIBILITY_DISPLAY_MAGNIFICATION_ENABLED, enabled ? ON : OFF);
+        Settings.Secure.putInt(mContext.getContentResolver(), TRIPLETAP_SHORTCUT_KEY,
+                enabled ? ON : OFF);
     }
 
     private String getStringFromSettings(String key) {
@@ -135,7 +233,57 @@ public class ToggleScreenMagnificationPreferenceFragmentTest {
     }
 
     private boolean getMagnificationTripleTapStatus() {
-        return Settings.Secure.getInt(mContext.getContentResolver(),
-                Settings.Secure.ACCESSIBILITY_DISPLAY_MAGNIFICATION_ENABLED, OFF) == ON;
+        return Settings.Secure.getInt(mContext.getContentResolver(), TRIPLETAP_SHORTCUT_KEY, OFF)
+                == ON;
+    }
+
+    private void callEmptyOnClicked(DialogInterface dialog, int which) {}
+
+    public static class TestToggleScreenMagnificationPreferenceFragment
+            extends ToggleScreenMagnificationPreferenceFragment {
+        @Override
+        protected void onPreferenceToggled(String preferenceKey, boolean enabled) {
+        }
+
+        @Override
+        public int getMetricsCategory() {
+            return 0;
+        }
+
+        @Override
+        int getUserShortcutTypes() {
+            return 0;
+        }
+
+        @Override
+        public int getPreferenceScreenResId() {
+            return R.xml.placeholder_prefs;
+        }
+
+        @Override
+        public View onCreateView(LayoutInflater inflater, ViewGroup container,
+                Bundle savedInstanceState) {
+            return mock(View.class);
+        }
+
+        @Override
+        public void onViewCreated(View view, Bundle savedInstanceState) {
+            // do nothing
+        }
+
+        @Override
+        public void onDestroyView() {
+            // do nothing
+        }
+
+        @Override
+        public void addPreferencesFromResource(@XmlRes int preferencesResId) {
+            // do nothing
+        }
+
+        @Override
+        protected void updateShortcutPreference() {
+            // UI related function, do nothing in tests
+        }
     }
 }