diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 497456f..a69f2ad 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -267,7 +267,7 @@
         val client = getKeyguardQuickAffordancePickerProviderClient(context)
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client, bgDispatcher),
+            KeyguardQuickAffordancePickerRepository(client, getFlags(), appContext),
             client
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
diff --git a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
index b17af80..080d70e 100644
--- a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
+++ b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
@@ -17,15 +17,14 @@
 
 package com.android.customization.picker.quickaffordance.data.repository
 
+import android.content.Context
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerAffordanceModel as AffordanceModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel as SelectionModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient as Client
-import com.android.systemui.shared.customization.data.content.CustomizationProviderContract as Contract
-import kotlinx.coroutines.CoroutineDispatcher
+import com.android.wallpaper.config.BaseFlags
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.withContext
 
 /**
  * Abstracts access to application state related to functionality for selecting, picking, or setting
@@ -33,11 +32,13 @@
  */
 class KeyguardQuickAffordancePickerRepository(
     private val client: Client,
-    private val backgroundDispatcher: CoroutineDispatcher,
+    private val flags: BaseFlags,
+    private val context: Context
 ) {
     /** Whether the feature is enabled. */
-    val isFeatureEnabled: Flow<Boolean> =
-        client.observeFlags().map { flags -> flags.isFeatureEnabled() }
+    fun isFeatureEnabled(): Boolean {
+        return flags.isQuickAffordancesEnabled(context)
+    }
 
     /** List of slots available on the device. */
     val slots: Flow<List<SlotModel>> =
@@ -55,18 +56,6 @@
             selections.map { selection -> selection.toModel() }
         }
 
-    suspend fun isFeatureEnabled(): Boolean {
-        return withContext(backgroundDispatcher) { client.queryFlags().isFeatureEnabled() }
-    }
-
-    private fun List<Client.Flag>.isFeatureEnabled(): Boolean {
-        return find { flag ->
-                flag.name ==
-                    Contract.FlagsTable.FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED
-            }
-            ?.value == true
-    }
-
     private fun Client.Slot.toModel(): SlotModel {
         return SlotModel(
             id = id,
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 35dbadd..38c43ca 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -17,19 +17,20 @@
 
 package com.android.customization.model.picker.quickaffordance.data.repository
 
+import android.content.Context
+import androidx.test.core.app.ApplicationProvider
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
-import com.android.systemui.shared.customization.data.content.CustomizationProviderContract
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
+import com.android.wallpaper.config.BaseFlags
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.toList
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
-import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
@@ -57,7 +58,15 @@
         underTest =
             KeyguardQuickAffordancePickerRepository(
                 client = client,
-                backgroundDispatcher = coroutineDispatcher,
+                flags =
+                    object : BaseFlags() {
+                        override fun getCachedFlags(
+                            context: Context
+                        ): List<CustomizationProviderClient.Flag> {
+                            return runBlocking { client.queryFlags() }
+                        }
+                    },
+                context = ApplicationProvider.getApplicationContext()
             )
     }
 
@@ -66,35 +75,9 @@
         Dispatchers.resetMain()
     }
 
+    // We need at least one test to prevent Studio errors
     @Test
-    fun `isFeatureEnabled - enabled`() =
-        testScope.runTest {
-            client.setFlag(
-                CustomizationProviderContract.FlagsTable
-                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
-                true,
-            )
-            val values = mutableListOf<Boolean>()
-            val job = launch { underTest.isFeatureEnabled.toList(values) }
-
-            assertThat(values.last()).isTrue()
-
-            job.cancel()
-        }
-
-    @Test
-    fun `isFeatureEnabled - not enabled`() =
-        testScope.runTest {
-            client.setFlag(
-                CustomizationProviderContract.FlagsTable
-                    .FLAG_NAME_CUSTOM_LOCK_SCREEN_QUICK_AFFORDANCES_ENABLED,
-                false,
-            )
-            val values = mutableListOf<Boolean>()
-            val job = launch { underTest.isFeatureEnabled.toList(values) }
-
-            assertThat(values.last()).isFalse()
-
-            job.cancel()
-        }
+    fun creationSucceeds() {
+        assertThat(underTest).isNotNull()
+    }
 }
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index efe9f64..2668064 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -17,6 +17,7 @@
 
 package com.android.customization.model.picker.quickaffordance.domain.interactor
 
+import androidx.test.core.app.ApplicationProvider
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
@@ -24,6 +25,7 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
@@ -57,12 +59,14 @@
         testScope = TestScope(testDispatcher)
         Dispatchers.setMain(testDispatcher)
         client = FakeCustomizationProviderClient()
+        val testFlags = object : BaseFlags() {}
         underTest =
             KeyguardQuickAffordancePickerInteractor(
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
-                        backgroundDispatcher = testDispatcher,
+                        flags = testFlags,
+                        context = ApplicationProvider.getApplicationContext(),
                     ),
                 client = client,
                 snapshotRestorer = {
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index f71bfc7..8dc8534 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -74,7 +74,8 @@
 
     @Before
     fun setUp() {
-        InjectorProvider.setInjector(TestInjector())
+        val injector = TestInjector()
+        InjectorProvider.setInjector(injector)
         context = ApplicationProvider.getApplicationContext()
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
@@ -86,7 +87,8 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
-                        backgroundDispatcher = testDispatcher,
+                        flags = injector.getFlags(),
+                        context = context,
                     ),
                 client = client,
                 snapshotRestorer = {