diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeModule.kt
index 0500a5844155f..8ae9e5e1fb8cd 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeModule.kt
@@ -215,9 +215,15 @@ abstract class ShadeModule {
         @Provides
         @SysUISingleton
         fun providesLockIconView(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            keyguardRootView: KeyguardRootView,
+            notificationPanelView: NotificationPanelView,
+            featureFlags: FeatureFlags
         ): LockIconView {
-            return notificationShadeWindowView.findViewById(R.id.lock_icon_view)
+            if (featureFlags.isEnabled(Flags.MIGRATE_LOCK_ICON)) {
+                return keyguardRootView.findViewById(R.id.lock_icon_view)
+            } else {
+                return notificationPanelView.findViewById(R.id.lock_icon_view)
+            }
         }
 
         // TODO(b/277762009): Only allow this view's controller to inject the view. See above.