diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 60a6803..84c84b0 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -58,6 +58,7 @@
         "Settings-robo-testutils",
         "Settings-testutils2",
         "SettingsLib-robo-testutils",
+        "SettingsLibPreference-testutils",
         "Settings_robolectric_meta_service_file",
         "aconfig_settings_flags_lib",
         "android.webkit.flags-aconfig-java",
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/FirmwareVersionScreenTest.kt b/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/FirmwareVersionScreenTest.kt
new file mode 100644
index 0000000..e0adc59
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/FirmwareVersionScreenTest.kt
@@ -0,0 +1,28 @@
+/*
+ * 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.deviceinfo.firmwareversion
+
+import com.android.settings.flags.Flags
+import com.android.settingslib.preference.CatalystScreenTestCase
+
+class FirmwareVersionScreenTest : CatalystScreenTestCase() {
+
+    override val preferenceScreenCreator = FirmwareVersionScreen()
+
+    override val flagName: String
+        get() = Flags.FLAG_CATALYST_FIRMWARE_VERSION
+}
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/MainlineModuleVersionPreferenceTest.kt b/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/MainlineModuleVersionPreferenceTest.kt
index 8fc1a2b..d1e8a8e 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/MainlineModuleVersionPreferenceTest.kt
+++ b/tests/robotests/src/com/android/settings/deviceinfo/firmwareversion/MainlineModuleVersionPreferenceTest.kt
@@ -150,7 +150,7 @@
         mainlineModuleVersionPreference.bind(preference, mainlineModuleVersionPreference)
 
         assertThat(preference.intent).isNull()
-        assertThat(preference.isSelectable).isFalse()
+        assertThat(preference.isSelectable).isTrue()
     }
 
     private fun createMocks(
diff --git a/tests/robotests/src/com/android/settings/deviceinfo/legal/LegalSettingsScreenTest.kt b/tests/robotests/src/com/android/settings/deviceinfo/legal/LegalSettingsScreenTest.kt
index 48e21b4..ce62844 100644
--- a/tests/robotests/src/com/android/settings/deviceinfo/legal/LegalSettingsScreenTest.kt
+++ b/tests/robotests/src/com/android/settings/deviceinfo/legal/LegalSettingsScreenTest.kt
@@ -15,39 +15,20 @@
  */
 package com.android.settings.deviceinfo.legal
 
-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.text.TextUtils
-import androidx.test.core.app.ApplicationProvider
 import com.android.settings.flags.Flags
+import com.android.settingslib.preference.CatalystScreenTestCase
 import com.google.common.truth.Truth.assertThat
-import org.junit.Rule
 import org.junit.Test
-import org.junit.runner.RunWith
-import org.robolectric.RobolectricTestRunner
 
-@RunWith(RobolectricTestRunner::class)
-class LegalSettingsScreenTest {
-    @get:Rule val setFlagsRule = SetFlagsRule()
-    private val context: Context = ApplicationProvider.getApplicationContext()
-    private val legalSettingsScreen = LegalSettingsScreen()
+class LegalSettingsScreenTest : CatalystScreenTestCase() {
+
+    override val preferenceScreenCreator = LegalSettingsScreen()
+
+    override val flagName: String
+        get() = Flags.FLAG_CATALYST_LEGAL_INFORMATION
 
     @Test
-    fun screenKey_exist() {
-        assertThat(TextUtils.equals(legalSettingsScreen.key, LegalSettingsScreen.KEY)).isTrue()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_CATALYST_LEGAL_INFORMATION)
-    fun isFlagEnabled_returnTrue() {
-        assertThat(legalSettingsScreen.isFlagEnabled(context)).isTrue()
-    }
-
-    @Test
-    @DisableFlags(Flags.FLAG_CATALYST_LEGAL_INFORMATION)
-    fun isFlagDisabled_returnTrue() {
-        assertThat(legalSettingsScreen.isFlagEnabled(context)).isFalse()
+    fun key() {
+        assertThat(preferenceScreenCreator.key).isEqualTo(LegalSettingsScreen.KEY)
     }
 }