Do not set Notetaking shortcut to default if Stylus is not detected

The current implementation set Notetaking shortcut as default in the keyguard for any device that contains Note Role - besides being able to manually add to the lock-screen, we should only set it as default whew a Stylus is detected (i.e., touch the screen or has been paired).

Testing instructions:
1. The device supports Note Task (i.e., tablets).
2. A Stylus.
3. Note Task feature flag is enabled: `adb shell cmd statusbar flag 1900 on`
4. Quick Affordance flag is enabled: `adb shell cmd statusbar flag 216 on`
5. Quick Affordance new UI flag is enabled: `adb shell cmd statusbar flag 222 on`

Test: atest SystemUITests:NoteTaskQuickAffordanceConfigTest

Fixes: b/265949213
Change-Id: I28d14ad2cc22b8d86434a6843df5358cfe9b089e

2 files changed, 145 insertions, 46 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfig.kt b/packages/SystemUI/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfig.kt
index 30660c492baa..8aed9958a7a2 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfig.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfig.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.notetask.quickaffordance
 
 import android.content.Context
+import android.hardware.input.InputSettings
 import com.android.systemui.R
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.ContentDescription
@@ -26,36 +27,52 @@ import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanc
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig.LockScreenState
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig.OnTriggeredResult
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig.PickerScreenState
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.notetask.NoteTaskController
 import com.android.systemui.notetask.NoteTaskEnabledKey
 import com.android.systemui.notetask.NoteTaskEntryPoint
+import com.android.systemui.stylus.StylusManager
+import dagger.Lazy
 import javax.inject.Inject
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.channels.trySendBlocking
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
 
 class NoteTaskQuickAffordanceConfig
 @Inject
 constructor(
     context: Context,
-    private val noteTaskController: NoteTaskController,
+    private val controller: NoteTaskController,
+    private val stylusManager: StylusManager,
+    private val lazyRepository: Lazy<KeyguardQuickAffordanceRepository>,
     @NoteTaskEnabledKey private val isEnabled: Boolean,
 ) : KeyguardQuickAffordanceConfig {
 
     override val key = BuiltInKeyguardQuickAffordanceKeys.CREATE_NOTE
 
-    override val pickerName: String = context.getString(R.string.note_task_button_label)
+    private val pickerNameResourceId = R.string.note_task_button_label
 
-    override val pickerIconResourceId = R.drawable.ic_note_task_shortcut_keyguard
+    override val pickerName: String = context.getString(pickerNameResourceId)
 
-    override val lockScreenState = flowOf(getLockScreenState())
+    override val pickerIconResourceId = R.drawable.ic_note_task_shortcut_keyguard
 
-    // TODO(b/265949213)
-    private fun getLockScreenState() =
-        if (isEnabled) {
-            val icon = Icon.Resource(pickerIconResourceId, ContentDescription.Loaded(pickerName))
-            LockScreenState.Visible(icon)
-        } else {
-            LockScreenState.Hidden
+    // Due to a dependency cycle with KeyguardQuickAffordanceRepository, we need to lazily access
+    // the repository when lockScreenState is accessed for the first time.
+    override val lockScreenState by lazy {
+        val stylusEverUsedFlow = createStylusEverUsedFlow(context, stylusManager)
+        val configSelectedFlow = createConfigSelectedFlow(lazyRepository.get(), key)
+        combine(configSelectedFlow, stylusEverUsedFlow) { isSelected, isStylusEverUsed ->
+            if (isEnabled && (isSelected || isStylusEverUsed)) {
+                val contentDescription = ContentDescription.Resource(pickerNameResourceId)
+                val icon = Icon.Resource(pickerIconResourceId, contentDescription)
+                LockScreenState.Visible(icon)
+            } else {
+                LockScreenState.Hidden
+            }
         }
+    }
 
     override suspend fun getPickerScreenState() =
         if (isEnabled) {
@@ -65,9 +82,27 @@ constructor(
         }
 
     override fun onTriggered(expandable: Expandable?): OnTriggeredResult {
-        noteTaskController.showNoteTask(
+        controller.showNoteTask(
             entryPoint = NoteTaskEntryPoint.QUICK_AFFORDANCE,
         )
         return OnTriggeredResult.Handled
     }
 }
+
+private fun createStylusEverUsedFlow(context: Context, stylusManager: StylusManager) =
+    callbackFlow {
+        trySendBlocking(InputSettings.isStylusEverUsed(context))
+        val callback =
+            object : StylusManager.StylusCallback {
+                override fun onStylusFirstUsed() {
+                    trySendBlocking(InputSettings.isStylusEverUsed(context))
+                }
+            }
+        stylusManager.registerCallback(callback)
+        awaitClose { stylusManager.unregisterCallback(callback) }
+    }
+
+private fun createConfigSelectedFlow(repository: KeyguardQuickAffordanceRepository, key: String) =
+    repository.selections.map { selected ->
+        selected.values.flatten().any { selectedConfig -> selectedConfig.key == key }
+    }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfigTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfigTest.kt
index e57d0d943aab..d44012f31fa5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/quickaffordance/NoteTaskQuickAffordanceConfigTest.kt
@@ -18,75 +18,139 @@
 
 package com.android.systemui.notetask.quickaffordance
 
+import android.hardware.input.InputSettings
 import android.test.suitebuilder.annotation.SmallTest
-import androidx.test.runner.AndroidJUnit4
+import android.testing.AndroidTestingRunner
+import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceConfig.LockScreenState
+import com.android.systemui.keyguard.data.repository.KeyguardQuickAffordanceRepository
 import com.android.systemui.notetask.NoteTaskController
 import com.android.systemui.notetask.NoteTaskEntryPoint
+import com.android.systemui.stylus.StylusManager
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runTest
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.verify
-import org.mockito.MockitoAnnotations
+import org.mockito.MockitoSession
+import org.mockito.quality.Strictness
 
-/**
- * Tests for [NoteTaskQuickAffordanceConfig].
- *
- * Build/Install/Run:
- * - atest SystemUITests:NoteTaskQuickAffordanceConfigTest
- */
+/** atest SystemUITests:NoteTaskQuickAffordanceConfigTest */
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(AndroidTestingRunner::class)
 internal class NoteTaskQuickAffordanceConfigTest : SysuiTestCase() {
 
-    @Mock lateinit var noteTaskController: NoteTaskController
+    @Mock lateinit var controller: NoteTaskController
+    @Mock lateinit var stylusManager: StylusManager
+    @Mock lateinit var repository: KeyguardQuickAffordanceRepository
+
+    private lateinit var mockitoSession: MockitoSession
 
     @Before
     fun setUp() {
-        MockitoAnnotations.initMocks(this)
+        mockitoSession =
+            ExtendedMockito.mockitoSession()
+                .initMocks(this)
+                .mockStatic(InputSettings::class.java)
+                .strictness(Strictness.LENIENT)
+                .startMocking()
+
+        whenever(InputSettings.isStylusEverUsed(mContext)).then { true }
+        whenever(repository.selections).then {
+            val map = mapOf("" to listOf(createUnderTest()))
+            MutableStateFlow(map)
+        }
     }
 
-    private fun createUnderTest(isEnabled: Boolean) =
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+    }
+
+    private fun createUnderTest(isEnabled: Boolean = true): KeyguardQuickAffordanceConfig =
         NoteTaskQuickAffordanceConfig(
             context = context,
-            noteTaskController = noteTaskController,
+            controller = controller,
+            stylusManager = stylusManager,
+            lazyRepository = { repository },
             isEnabled = isEnabled,
         )
 
+    private fun createLockScreenStateVisible(): LockScreenState =
+        LockScreenState.Visible(
+            icon =
+                Icon.Resource(
+                    res = R.drawable.ic_note_task_shortcut_keyguard,
+                    contentDescription =
+                        ContentDescription.Resource(R.string.note_task_button_label),
+                )
+        )
+
     @Test
-    fun lockScreenState_isNotEnabled_shouldEmitHidden() = runTest {
-        val underTest = createUnderTest(isEnabled = false)
+    fun lockScreenState_stylusUsed_noCustomShortcutSelected_shouldEmitVisible() = runTest {
+        val underTest = createUnderTest()
+
+        val actual by collectLastValue(underTest.lockScreenState)
+
+        assertThat(actual).isEqualTo(createLockScreenStateVisible())
+    }
+
+    @Test
+    fun lockScreenState_noStylusEverUsed_noCustomShortcutSelected_shouldEmitVisible() = runTest {
+        whenever(InputSettings.isStylusEverUsed(mContext)).then { false }
+        val underTest = createUnderTest()
 
-        val actual = collectLastValue(underTest.lockScreenState)
+        val actual by collectLastValue(underTest.lockScreenState)
 
-        assertThat(actual()).isEqualTo(LockScreenState.Hidden)
+        assertThat(actual).isEqualTo(createLockScreenStateVisible())
     }
 
     @Test
-    fun lockScreenState_isEnabled_shouldEmitVisible() = runTest {
-        val stringResult = "Notetaking"
-        val underTest = createUnderTest(isEnabled = true)
-
-        val actual = collectLastValue(underTest.lockScreenState)
-
-        val expected =
-            LockScreenState.Visible(
-                icon =
-                    Icon.Resource(
-                        res = R.drawable.ic_note_task_shortcut_keyguard,
-                        contentDescription = ContentDescription.Loaded(stringResult),
-                    )
-            )
-        assertThat(actual()).isEqualTo(expected)
+    fun lockScreenState_stylusUsed_customShortcutSelected_shouldEmitVisible() = runTest {
+        whenever(repository.selections).then {
+            val map = mapOf<String, List<KeyguardQuickAffordanceConfig>>()
+            MutableStateFlow(map)
+        }
+        val underTest = createUnderTest()
+
+        val actual by collectLastValue(underTest.lockScreenState)
+
+        assertThat(actual).isEqualTo(createLockScreenStateVisible())
+    }
+
+    @Test
+    fun lockScreenState_noIsStylusEverUsed_noCustomShortcutSelected_shouldEmitHidden() = runTest {
+        whenever(InputSettings.isStylusEverUsed(mContext)).then { false }
+        whenever(repository.selections).then {
+            val map = mapOf<String, List<KeyguardQuickAffordanceConfig>>()
+            MutableStateFlow(map)
+        }
+        val underTest = createUnderTest()
+
+        val actual by collectLastValue(underTest.lockScreenState)
+
+        assertThat(actual).isEqualTo(LockScreenState.Hidden)
+    }
+
+    @Test
+    fun lockScreenState_isNotEnabled_shouldEmitHidden() = runTest {
+        val underTest = createUnderTest(isEnabled = false)
+
+        val actual by collectLastValue(underTest.lockScreenState)
+
+        assertThat(actual).isEqualTo(LockScreenState.Hidden)
     }
 
     @Test
@@ -95,6 +159,6 @@ internal class NoteTaskQuickAffordanceConfigTest : SysuiTestCase() {
 
         underTest.onTriggered(expandable = null)
 
-        verify(noteTaskController).showNoteTask(entryPoint = NoteTaskEntryPoint.QUICK_AFFORDANCE)
+        verify(controller).showNoteTask(entryPoint = NoteTaskEntryPoint.QUICK_AFFORDANCE)
     }
 }