diff --git a/src/com/android/settings/network/AllInOneTetherPreferenceController.java b/src/com/android/settings/network/AllInOneTetherPreferenceController.java
index 8180d5a..75453b4 100644
--- a/src/com/android/settings/network/AllInOneTetherPreferenceController.java
+++ b/src/com/android/settings/network/AllInOneTetherPreferenceController.java
@@ -37,6 +37,7 @@
 import androidx.lifecycle.Lifecycle.Event;
 import androidx.lifecycle.LifecycleObserver;
 import androidx.lifecycle.OnLifecycleEvent;
+import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
 import com.android.settings.core.BasePreferenceController;
@@ -94,6 +95,12 @@
     }
 
     @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = screen.findPreference(mPreferenceKey);
+    }
+
+    @Override
     public int getAvailabilityStatus() {
         if (!TetherUtil.isTetherAvailable(mContext)
                 || !FeatureFlagUtils.isEnabled(mContext, FeatureFlags.TETHER_ALL_IN_ONE)) {
diff --git a/tests/robotests/src/com/android/settings/network/AllInOneTetherPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/network/AllInOneTetherPreferenceControllerTest.java
index 12e687d..f37b0e4 100644
--- a/tests/robotests/src/com/android/settings/network/AllInOneTetherPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/network/AllInOneTetherPreferenceControllerTest.java
@@ -24,6 +24,7 @@
 
 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;
@@ -35,6 +36,9 @@
 import android.bluetooth.BluetoothProfile;
 import android.content.Context;
 
+import androidx.preference.PreferenceScreen;
+import androidx.test.core.app.ApplicationProvider;
+
 import com.android.settings.R;
 import com.android.settings.widget.MasterSwitchPreference;
 
@@ -109,13 +113,13 @@
         });
     }
 
-    @Mock
     private Context mContext;
     @Mock
     private BluetoothAdapter mBluetoothAdapter;
     @Mock
     private MasterSwitchPreference mPreference;
 
+    private static final String PREF_KEY = "tether";
     private AllInOneTetherPreferenceController mController;
     private final int mTetherState;
     private final int mSummaryResId;
@@ -127,11 +131,16 @@
 
     @Before
     public void setUp() {
+        mContext = ApplicationProvider.getApplicationContext();
         MockitoAnnotations.initMocks(this);
         mController = spy(AllInOneTetherPreferenceController.class);
         ReflectionHelpers.setField(mController, "mContext", mContext);
         ReflectionHelpers.setField(mController, "mBluetoothAdapter", mBluetoothAdapter);
-        ReflectionHelpers.setField(mController, "mPreference", mPreference);
+        ReflectionHelpers.setField(mController, "mPreferenceKey", PREF_KEY);
+        PreferenceScreen screen = mock(PreferenceScreen.class);
+        when(screen.findPreference(PREF_KEY)).thenReturn(mPreference);
+        doReturn(mController.AVAILABLE).when(mController).getAvailabilityStatus();
+        mController.displayPreference(screen);
     }
 
     @Test
@@ -169,5 +178,8 @@
     public void getSummary_afterTetherStateChanged() {
         mController.onTetherStateUpdated(mTetherState);
         assertThat(mController.getSummary()).isEqualTo(mContext.getString(mSummaryResId));
+
+        verify(mController).updateState(mPreference);
+        verify(mPreference).setSummary(mContext.getString(mSummaryResId));
     }
 }