diff --git a/packages/SettingsLib/tests/robotests/Android.bp b/packages/SettingsLib/tests/robotests/Android.bp
index 75c40bf..f380e7f 100644
--- a/packages/SettingsLib/tests/robotests/Android.bp
+++ b/packages/SettingsLib/tests/robotests/Android.bp
@@ -56,6 +56,7 @@
         "flag-junit",
         "settingslib_media_flags_lib",
         "settingslib_illustrationpreference_flags_lib",
+        "settingslib_selectorwithwidgetpreference_flags_lib",
         "testng", // TODO: remove once JUnit on Android provides assertThrows
     ],
     java_resource_dirs: ["config"],
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/SelectorWithWidgetPreferenceTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/SelectorWithWidgetPreferenceTest.java
index 60885f1..243ce85 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/SelectorWithWidgetPreferenceTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/SelectorWithWidgetPreferenceTest.java
@@ -23,22 +23,31 @@
 import static org.mockito.Mockito.when;
 
 import android.app.Application;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
+import android.util.AttributeSet;
 import android.view.LayoutInflater;
 import android.view.View;
+import android.widget.TextView;
 
 import androidx.preference.PreferenceViewHolder;
+import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settingslib.widget.preference.selector.R;
+import com.android.settingslib.widget.selectorwithwidgetpreference.flags.Flags;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 @RunWith(RobolectricTestRunner.class)
 public class SelectorWithWidgetPreferenceTest {
 
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
     private Application mContext;
     private SelectorWithWidgetPreference mPreference;
 
@@ -55,7 +64,7 @@
 
     @Before
     public void setUp() {
-        mContext = RuntimeEnvironment.application;
+        mContext = ApplicationProvider.getApplicationContext();
         mPreference = new SelectorWithWidgetPreference(mContext);
 
         View view = LayoutInflater.from(mContext)
@@ -121,6 +130,60 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ALLOW_SET_TITLE_MAX_LINES)
+    public void onBindViewHolder_titleMaxLinesSet_flagOff_titleMaxLinesMatchesDefault() {
+        final int titleMaxLines = 5;
+        AttributeSet attributeSet = Robolectric.buildAttributeSet()
+                .addAttribute(R.attr.titleMaxLines, String.valueOf(titleMaxLines))
+                .build();
+        mPreference = new SelectorWithWidgetPreference(mContext, attributeSet);
+        View view = LayoutInflater.from(mContext)
+                .inflate(mPreference.getLayoutResource(), null /* root */);
+        PreferenceViewHolder preferenceViewHolder =
+                PreferenceViewHolder.createInstanceForTests(view);
+
+        mPreference.onBindViewHolder(preferenceViewHolder);
+
+        TextView title = (TextView) preferenceViewHolder.findViewById(android.R.id.title);
+        assertThat(title.getMaxLines()).isEqualTo(SelectorWithWidgetPreference.DEFAULT_MAX_LINES);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ALLOW_SET_TITLE_MAX_LINES)
+    public void onBindViewHolder_noTitleMaxLinesSet_titleMaxLinesMatchesDefault() {
+        AttributeSet attributeSet = Robolectric.buildAttributeSet().build();
+        mPreference = new SelectorWithWidgetPreference(mContext, attributeSet);
+        View view = LayoutInflater.from(mContext)
+                .inflate(mPreference.getLayoutResource(), null /* root */);
+        PreferenceViewHolder preferenceViewHolder =
+                PreferenceViewHolder.createInstanceForTests(view);
+
+        mPreference.onBindViewHolder(preferenceViewHolder);
+
+        TextView title = (TextView) preferenceViewHolder.findViewById(android.R.id.title);
+        assertThat(title.getMaxLines()).isEqualTo(SelectorWithWidgetPreference.DEFAULT_MAX_LINES);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ALLOW_SET_TITLE_MAX_LINES)
+    public void onBindViewHolder_titleMaxLinesSet_titleMaxLinesUpdated() {
+        final int titleMaxLines = 5;
+        AttributeSet attributeSet = Robolectric.buildAttributeSet()
+                .addAttribute(R.attr.titleMaxLines, String.valueOf(titleMaxLines))
+                .build();
+        mPreference = new SelectorWithWidgetPreference(mContext, attributeSet);
+        View view = LayoutInflater.from(mContext)
+                .inflate(mPreference.getLayoutResource(), null /* root */);
+        PreferenceViewHolder preferenceViewHolder =
+                PreferenceViewHolder.createInstanceForTests(view);
+
+        mPreference.onBindViewHolder(preferenceViewHolder);
+
+        TextView title = (TextView) preferenceViewHolder.findViewById(android.R.id.title);
+        assertThat(title.getMaxLines()).isEqualTo(titleMaxLines);
+    }
+
+    @Test
     public void nullSummary_containerShouldBeGone() {
         mPreference.setSummary(null);
         View summaryContainer = new View(mContext);