[TP] Clock Registry Provider

Create ClockRegistryProvider since we need the plugin manager to be
connected to use ClockRegistry

Test: Manually test that data can be fetched correctly from the registry
Bug: 262924055
Change-Id: I77883044224ff2639b7a96b0d1c7d94db9deb8ca
diff --git a/res/layout/clock_section_view.xml b/res/layout/clock_section_view.xml
index 4a651ca..dfd015e 100644
--- a/res/layout/clock_section_view.xml
+++ b/res/layout/clock_section_view.xml
@@ -14,7 +14,7 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.customization.picker.clock.ClockSectionView
+<com.android.customization.picker.clock.ui.view.ClockSectionView
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
@@ -51,4 +51,4 @@
         android:src="@drawable/ic_clock_24px"
         android:background="@drawable/option_border_color"
         android:contentDescription="@string/clock_picker_entry_content_description" />
-</com.android.customization.picker.clock.ClockSectionView>
\ No newline at end of file
+</com.android.customization.picker.clock.ui.view.ClockSectionView>
\ No newline at end of file
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
index f5d6c54..6194e11 100644
--- a/src/com/android/customization/module/CustomizationInjector.kt
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -20,9 +20,10 @@
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
+import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
+import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
-import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.wallpaper.module.Injector
 
@@ -40,9 +41,15 @@
         context: Context
     ): KeyguardQuickAffordancePickerInteractor
 
-    fun getClockRegistry(context: Context): ClockRegistry
+    fun getClockRegistryProvider(context: Context): ClockRegistryProvider
 
-    fun getClockSectionViewModel(context: Context): ClockSectionViewModel
+    fun getClockPickerInteractor(
+        context: Context,
+        clockRegistry: ClockRegistry
+    ): ClockPickerInteractor
 
-    fun getPluginManager(context: Context): PluginManager
+    fun getClockSectionViewModel(
+        context: Context,
+        clockRegistry: ClockRegistry
+    ): ClockSectionViewModel
 }
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 5239b75..d80086e 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -15,15 +15,10 @@
  */
 package com.android.customization.module
 
-import android.app.NotificationManager
-import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
 import android.net.Uri
 import android.os.Bundle
-import android.os.Handler
-import android.os.UserHandle
-import android.view.LayoutInflater
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
@@ -31,6 +26,7 @@
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
+import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.notifications.data.repository.NotificationsRepository
@@ -40,18 +36,9 @@
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
-import com.android.systemui.plugins.Plugin
-import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.systemui.shared.clocks.DefaultClockProvider
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
-import com.android.systemui.shared.plugins.PluginActionManager
-import com.android.systemui.shared.plugins.PluginEnabler
-import com.android.systemui.shared.plugins.PluginInstance
-import com.android.systemui.shared.plugins.PluginManagerImpl
-import com.android.systemui.shared.plugins.PluginPrefs
-import com.android.systemui.shared.system.UncaughtExceptionPreHandlerManager_Factory
 import com.android.wallpaper.model.LiveWallpaperInfo
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.CustomizationSections
@@ -64,7 +51,6 @@
 import com.android.wallpaper.picker.LivePreviewFragment
 import com.android.wallpaper.picker.PreviewFragment
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
-import java.util.concurrent.Executors
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 
@@ -81,9 +67,9 @@
     private var fragmentFactory: FragmentFactory? = null
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
-    private var clockRegistry: ClockRegistry? = null
+    private var clockRegistryProvider: ClockRegistryProvider? = null
+    private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
-    private var pluginManager: PluginManager? = null
     private var notificationsInteractor: NotificationsInteractor? = null
     private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
 
@@ -234,86 +220,29 @@
                 .also { keyguardQuickAffordanceSnapshotRestorer = it }
     }
 
-    override fun getClockRegistry(context: Context): ClockRegistry {
-        return clockRegistry
-            ?: ClockRegistry(
-                    context,
-                    getPluginManager(context),
-                    Handler.getMain(),
-                    isEnabled = true,
-                    userHandle = UserHandle.USER_SYSTEM,
-                    DefaultClockProvider(context, LayoutInflater.from(context), context.resources)
-                )
-                .also { clockRegistry = it }
+    override fun getClockRegistryProvider(context: Context): ClockRegistryProvider {
+        return clockRegistryProvider
+            ?: ClockRegistryProvider(context).also { clockRegistryProvider = it }
     }
 
-    override fun getClockSectionViewModel(context: Context): ClockSectionViewModel {
-        return clockSectionViewModel
-            ?: ClockSectionViewModel(
-                    ClockPickerInteractor(ClockPickerRepositoryImpl(getClockRegistry(context)))
-                )
-                .also { clockSectionViewModel = it }
-    }
-
-    override fun getPluginManager(context: Context): PluginManager {
-        return pluginManager ?: createPluginManager(context).also { pluginManager = it }
-    }
-
-    private fun createPluginManager(context: Context): PluginManager {
-        val privilegedPlugins = listOf<String>()
-        val isDebugDevice = true
-
-        val instanceFactory =
-            PluginInstance.Factory(
-                this::class.java.classLoader,
-                PluginInstance.InstanceFactory<Plugin>(),
-                PluginInstance.VersionChecker(),
-                privilegedPlugins,
-                isDebugDevice,
-            )
-
-        /*
-         * let SystemUI handle plugin, in this class assume plugins are enabled
-         */
-        val pluginEnabler =
-            object : PluginEnabler {
-                override fun setEnabled(component: ComponentName) = Unit
-
-                override fun setDisabled(
-                    component: ComponentName,
-                    @PluginEnabler.DisableReason reason: Int
-                ) = Unit
-
-                override fun isEnabled(component: ComponentName): Boolean {
-                    return true
-                }
-
-                @PluginEnabler.DisableReason
-                override fun getDisableReason(componentName: ComponentName): Int {
-                    return PluginEnabler.ENABLED
-                }
+    override fun getClockPickerInteractor(
+        context: Context,
+        clockRegistry: ClockRegistry,
+    ): ClockPickerInteractor {
+        return clockPickerInteractor
+            ?: ClockPickerInteractor(ClockPickerRepositoryImpl(clockRegistry)).also {
+                clockPickerInteractor = it
             }
+    }
 
-        val pluginActionManager =
-            PluginActionManager.Factory(
-                context,
-                context.packageManager,
-                context.mainExecutor,
-                Executors.newSingleThreadExecutor(),
-                context.getSystemService(NotificationManager::class.java),
-                pluginEnabler,
-                privilegedPlugins,
-                instanceFactory,
-            )
-        return PluginManagerImpl(
-            context,
-            pluginActionManager,
-            isDebugDevice,
-            UncaughtExceptionPreHandlerManager_Factory.create().get(),
-            pluginEnabler,
-            PluginPrefs(context),
-            listOf(),
-        )
+    override fun getClockSectionViewModel(
+        context: Context,
+        clockRegistry: ClockRegistry,
+    ): ClockSectionViewModel {
+        return clockSectionViewModel
+            ?: ClockSectionViewModel(getClockPickerInteractor(context, clockRegistry)).also {
+                clockSectionViewModel = it
+            }
     }
 
     protected fun getNotificationsInteractor(
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
new file mode 100644
index 0000000..2529958
--- /dev/null
+++ b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2023 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.customization.picker.clock.data.repository
+
+import android.app.NotificationManager
+import android.content.ComponentName
+import android.content.Context
+import android.os.Handler
+import android.os.UserHandle
+import android.view.LayoutInflater
+import com.android.systemui.plugins.ClockProviderPlugin
+import com.android.systemui.plugins.Plugin
+import com.android.systemui.plugins.PluginListener
+import com.android.systemui.plugins.PluginManager
+import com.android.systemui.shared.clocks.ClockRegistry
+import com.android.systemui.shared.clocks.DefaultClockProvider
+import com.android.systemui.shared.plugins.PluginActionManager
+import com.android.systemui.shared.plugins.PluginEnabler
+import com.android.systemui.shared.plugins.PluginInstance
+import com.android.systemui.shared.plugins.PluginManagerImpl
+import com.android.systemui.shared.plugins.PluginPrefs
+import com.android.systemui.shared.system.UncaughtExceptionPreHandlerManager_Factory
+import java.util.concurrent.Executors
+import kotlinx.coroutines.suspendCancellableCoroutine
+
+/**
+ * Provide the [ClockRegistry] singleton. Note that we need to make sure that the [PluginManager]
+ * needs to be connected before [ClockRegistry] is ready to use.
+ */
+class ClockRegistryProvider(
+    private val context: Context,
+) {
+    private val pluginManager: PluginManager by lazy { createPluginManager(context) }
+    private val clockRegistry: ClockRegistry by lazy {
+        ClockRegistry(
+            context,
+            pluginManager,
+            Handler.getMain(),
+            isEnabled = true,
+            userHandle = UserHandle.USER_SYSTEM,
+            DefaultClockProvider(context, LayoutInflater.from(context), context.resources)
+        )
+    }
+
+    suspend fun get(): ClockRegistry {
+        return suspendCancellableCoroutine { continuation ->
+            val pluginListener =
+                object : PluginListener<ClockProviderPlugin> {
+                    var hasConnected = false
+                    override fun onPluginConnected(
+                        plugin: ClockProviderPlugin?,
+                        pluginContext: Context?
+                    ) {
+                        if (!hasConnected) {
+                            pluginManager.removePluginListener(this)
+                            hasConnected = true
+                            continuation.resumeWith(Result.success(clockRegistry))
+                        }
+                    }
+                }
+            pluginManager.addPluginListener(pluginListener, ClockProviderPlugin::class.java, true)
+            continuation.invokeOnCancellation { pluginManager.removePluginListener(pluginListener) }
+        }
+    }
+
+    private fun createPluginManager(context: Context): PluginManager {
+        val privilegedPlugins = listOf<String>()
+        val isDebugDevice = true
+
+        val instanceFactory =
+            PluginInstance.Factory(
+                this::class.java.classLoader,
+                PluginInstance.InstanceFactory<Plugin>(),
+                PluginInstance.VersionChecker(),
+                privilegedPlugins,
+                isDebugDevice,
+            )
+
+        /*
+         * let SystemUI handle plugin, in this class assume plugins are enabled
+         */
+        val pluginEnabler =
+            object : PluginEnabler {
+                override fun setEnabled(component: ComponentName) = Unit
+
+                override fun setDisabled(
+                    component: ComponentName,
+                    @PluginEnabler.DisableReason reason: Int
+                ) = Unit
+
+                override fun isEnabled(component: ComponentName): Boolean {
+                    return true
+                }
+
+                @PluginEnabler.DisableReason
+                override fun getDisableReason(componentName: ComponentName): Int {
+                    return PluginEnabler.ENABLED
+                }
+            }
+
+        val pluginActionManager =
+            PluginActionManager.Factory(
+                context,
+                context.packageManager,
+                context.mainExecutor,
+                Executors.newSingleThreadExecutor(),
+                context.getSystemService(NotificationManager::class.java),
+                pluginEnabler,
+                privilegedPlugins,
+                instanceFactory,
+            )
+        return PluginManagerImpl(
+            context,
+            pluginActionManager,
+            isDebugDevice,
+            UncaughtExceptionPreHandlerManager_Factory.create().get(),
+            pluginEnabler,
+            PluginPrefs(context),
+            listOf(),
+        )
+    }
+}
diff --git a/src/com/android/customization/picker/clock/ClockCustomDemoFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
similarity index 73%
rename from src/com/android/customization/picker/clock/ClockCustomDemoFragment.kt
rename to src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
index 4f65080..4aa5de4 100644
--- a/src/com/android/customization/picker/clock/ClockCustomDemoFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
@@ -1,4 +1,4 @@
-package com.android.customization.picker.clock
+package com.android.customization.picker.clock.ui.fragment
 
 import android.content.Context
 import android.os.Bundle
@@ -13,43 +13,23 @@
 import android.widget.TextView
 import android.widget.Toast
 import androidx.core.view.setPadding
+import androidx.lifecycle.lifecycleScope
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.module.ThemePickerInjector
 import com.android.internal.annotations.VisibleForTesting
 import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.plugins.ClockProviderPlugin
-import com.android.systemui.plugins.PluginListener
-import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 class ClockCustomDemoFragment : AppbarFragment() {
-    @VisibleForTesting lateinit var clockRegistry: ClockRegistry
     @VisibleForTesting lateinit var recyclerView: RecyclerView
-    lateinit var pluginManager: PluginManager
-    @VisibleForTesting
-    val pluginListener =
-        object : PluginListener<ClockProviderPlugin> {
-            override fun onPluginConnected(plugin: ClockProviderPlugin, context: Context) {
-                val listInUse = clockRegistry.getClocks().filter { "NOT_IN_USE" !in it.clockId }
-                recyclerView.adapter =
-                    ClockRecyclerAdapter(listInUse, context) {
-                        clockRegistry.currentClockId = it.clockId
-                        Toast.makeText(context, "${it.name} selected", Toast.LENGTH_SHORT).show()
-                    }
-            }
-        }
-
-    override fun onAttach(context: Context) {
-        super.onAttach(context)
-        val injector = InjectorProvider.getInjector() as ThemePickerInjector
-        pluginManager = injector.getPluginManager(context)
-        clockRegistry = injector.getClockRegistry(context)
-        pluginManager.addPluginListener(pluginListener, ClockProviderPlugin::class.java, true)
-    }
+    @VisibleForTesting lateinit var clockRegistry: ClockRegistry
 
     override fun onCreateView(
         inflater: LayoutInflater,
@@ -58,6 +38,20 @@
     ): View {
         val view = inflater.inflate(R.layout.fragment_clock_custom_picker_demo, container, false)
         setUpToolbar(view)
+        lifecycleScope.launch {
+            clockRegistry =
+                withContext(Dispatchers.IO) {
+                    (InjectorProvider.getInjector() as ThemePickerInjector)
+                        .getClockRegistryProvider(requireContext())
+                        .get()
+                }
+            val listInUse = clockRegistry.getClocks().filter { "NOT_IN_USE" !in it.clockId }
+            recyclerView.adapter =
+                ClockRecyclerAdapter(listInUse, requireContext()) {
+                    clockRegistry.currentClockId = it.clockId
+                    Toast.makeText(context, "${it.name} selected", Toast.LENGTH_SHORT).show()
+                }
+        }
         return view
     }
 
diff --git a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt b/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
index 748fa17..052d0f5 100644
--- a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
+++ b/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
@@ -18,21 +18,27 @@
 import android.content.Context
 import android.view.LayoutInflater
 import androidx.lifecycle.LifecycleOwner
-import com.android.customization.picker.clock.ClockCustomDemoFragment
-import com.android.customization.picker.clock.ClockSectionView
+import androidx.lifecycle.lifecycleScope
+import com.android.customization.module.ThemePickerInjector
+import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.ui.binder.ClockSectionViewBinder
-import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
+import com.android.customization.picker.clock.ui.fragment.ClockCustomDemoFragment
+import com.android.customization.picker.clock.ui.view.ClockSectionView
 import com.android.wallpaper.R
 import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.model.CustomizationSectionController
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
+import com.android.wallpaper.module.InjectorProvider
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 /** A [CustomizationSectionController] for clock customization. */
 class ClockSectionController(
     private val navigationController: CustomizationSectionNavigationController,
-    private val viewModel: ClockSectionViewModel,
     private val lifecycleOwner: LifecycleOwner,
     private val flag: BaseFlags,
+    private val clockRegistryProvider: ClockRegistryProvider,
 ) : CustomizationSectionController<ClockSectionView?> {
 
     override fun isAvailable(context: Context?): Boolean {
@@ -46,11 +52,16 @@
                     R.layout.clock_section_view,
                     null,
                 ) as ClockSectionView
-        ClockSectionViewBinder.bind(
-            view = view,
-            viewModel = viewModel,
-            lifecycleOwner = lifecycleOwner
-        ) { navigationController.navigateTo(ClockCustomDemoFragment()) }
+        lifecycleOwner.lifecycleScope.launch {
+            val registry = withContext(Dispatchers.IO) { clockRegistryProvider.get() }
+            ClockSectionViewBinder.bind(
+                view = view,
+                viewModel =
+                    (InjectorProvider.getInjector() as ThemePickerInjector)
+                        .getClockSectionViewModel(context, registry),
+                lifecycleOwner = lifecycleOwner
+            ) { navigationController.navigateTo(ClockCustomDemoFragment()) }
+        }
         return view
     }
 }
diff --git a/src/com/android/customization/picker/clock/ClockSectionView.kt b/src/com/android/customization/picker/clock/ui/view/ClockSectionView.kt
similarity index 93%
rename from src/com/android/customization/picker/clock/ClockSectionView.kt
rename to src/com/android/customization/picker/clock/ui/view/ClockSectionView.kt
index fac975a..cca107c 100644
--- a/src/com/android/customization/picker/clock/ClockSectionView.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockSectionView.kt
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.customization.picker.clock
+package com.android.customization.picker.clock.ui.view
 
 import android.content.Context
 import android.util.AttributeSet
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ClockCustomDemoFragmentTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
similarity index 93%
rename from tests/robotests/src/com/android/customization/picker/clock/ClockCustomDemoFragmentTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
index ad3dd1c..430dd27 100644
--- a/tests/robotests/src/com/android/customization/picker/clock/ClockCustomDemoFragmentTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
@@ -1,4 +1,4 @@
-package com.android.customization.picker.clock
+package com.android.customization.picker.clock.ui.fragment
 
 import android.os.Handler
 import android.os.UserHandle
@@ -9,7 +9,6 @@
 import com.android.systemui.plugins.ClockId
 import com.android.systemui.plugins.ClockMetadata
 import com.android.systemui.plugins.ClockProvider
-import com.android.systemui.plugins.ClockProviderPlugin
 import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
 import org.junit.Assert
@@ -32,7 +31,6 @@
     private lateinit var registry: ClockRegistry
     @Mock private lateinit var mockPluginManager: PluginManager
     @Mock private lateinit var mockHandler: Handler
-    @Mock private lateinit var fakePlugin: ClockProviderPlugin
     @Mock private lateinit var defaultClockProvider: ClockProvider
 
     private var settingValue: String = ""
@@ -74,7 +72,6 @@
         mClockCustomDemoFragment!!.recyclerView = RecyclerView(mActivity)
         mClockCustomDemoFragment!!.recyclerView.layoutManager =
             LinearLayoutManager(mActivity, RecyclerView.VERTICAL, false)
-        mClockCustomDemoFragment!!.pluginListener.onPluginConnected(fakePlugin, mActivity)
     }
 
     @Test
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
index eb135d9..735ba8e 100644
--- a/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -1,9 +1,6 @@
 package com.android.customization.testing
 
 import android.content.Context
-import android.os.Handler
-import android.os.UserHandle
-import android.view.LayoutInflater
 import androidx.fragment.app.FragmentActivity
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
@@ -11,15 +8,14 @@
 import com.android.customization.module.CustomizationInjector
 import com.android.customization.module.CustomizationPreferences
 import com.android.customization.module.ThemesUserEventLogger
+import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.data.repository.FakeClockPickerRepository
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
-import com.android.systemui.plugins.PluginManager
 import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.systemui.shared.clocks.DefaultClockProvider
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
 import com.android.wallpaper.config.BaseFlags
@@ -43,9 +39,9 @@
     private var customizationProviderClient: CustomizationProviderClient? = null
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
-    private var clockRegistry: ClockRegistry? = null
+    private var clockRegistryProvider: ClockRegistryProvider? = null
+    private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
-    private var pluginManager: PluginManager? = null
 
     override fun getCustomizationPreferences(context: Context): CustomizationPreferences {
         return customizationPreferences
@@ -131,28 +127,29 @@
                 .also { keyguardQuickAffordanceSnapshotRestorer = it }
     }
 
-    override fun getClockRegistry(context: Context): ClockRegistry {
-        return clockRegistry
-            ?: ClockRegistry(
-                    context,
-                    getPluginManager(context),
-                    Handler.getMain(),
-                    isEnabled = true,
-                    userHandle = UserHandle.USER_SYSTEM,
-                    DefaultClockProvider(context, LayoutInflater.from(context), context.resources)
-                )
-                .also { clockRegistry = it }
+    override fun getClockRegistryProvider(context: Context): ClockRegistryProvider {
+        return clockRegistryProvider
+            ?: ClockRegistryProvider(context).also { clockRegistryProvider = it }
     }
 
-    override fun getClockSectionViewModel(context: Context): ClockSectionViewModel {
-        return clockSectionViewModel
-            ?: ClockSectionViewModel(ClockPickerInteractor(FakeClockPickerRepository())).also {
-                clockSectionViewModel = it
+    override fun getClockPickerInteractor(
+        context: Context,
+        clockRegistry: ClockRegistry
+    ): ClockPickerInteractor {
+        return clockPickerInteractor
+            ?: ClockPickerInteractor(FakeClockPickerRepository()).also {
+                clockPickerInteractor = it
             }
     }
 
-    override fun getPluginManager(context: Context): PluginManager {
-        return pluginManager ?: TestPluginManager().also { pluginManager = it }
+    override fun getClockSectionViewModel(
+        context: Context,
+        clockRegistry: ClockRegistry
+    ): ClockSectionViewModel {
+        return clockSectionViewModel
+            ?: ClockSectionViewModel(getClockPickerInteractor(context, clockRegistry)).also {
+                clockSectionViewModel = it
+            }
     }
 
     companion object {