Fix pinning option semantics (#3668)

Author: rli

.changes/next-release/bugfix-fa693f13-dfc3-4fd3-9116-b8b4f0a7d048.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix issue where the \"Do not ask again\" option is not respected when switching connections on CodeWhisperer/CodeCatalyst"
+}

jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManager.kt

@@ -56,7 +56,9 @@ class DefaultConnectionPinningManager :
     ConnectionPinningManager,
     PersistentStateComponent<ConnectionPinningManagerState>,
     Disposable {
-    private var doNotPromptForPinning: Boolean = false
+    @VisibleForTesting
+    internal var shouldPinConnections: Boolean? = null
+
     private val pinnedConnections = ConcurrentHashMap<String, ToolkitConnection>()
 
     init {
@@ -108,14 +110,14 @@ class DefaultConnectionPinningManager :
     }
 
     override fun getState() = ConnectionPinningManagerState(
-        doNotPromptForPinning,
+        shouldPinConnections,
         pinnedConnections.entries.associate { (k, v) -> k to v.id }
     )
 
     override fun loadState(state: ConnectionPinningManagerState) {
         val authManager = ToolkitAuthManager.getInstance()
 
-        doNotPromptForPinning = state.doNotPromptForPinning
+        shouldPinConnections = state.shouldPinConnections
 
         pinnedConnections.clear()
         pinnedConnections.putAll(
@@ -129,35 +131,37 @@ class DefaultConnectionPinningManager :
 
     @VisibleForTesting
     internal fun showDialogIfNeeded(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection, featuresString: String, project: Project? = null) =
-        if (!doNotPromptForPinning) {
-            val bearerTokenConnectionName = bearerTokenConnectionString(oldConnection, newConnection)
-
-            computeOnEdt(ModalityState.defaultModalityState()) {
-                MessageDialogBuilder.yesNo(
-                    message("credentials.switch.confirmation.title", featuresString, bearerTokenConnectionName),
-                    message("credentials.switch.confirmation.comment", featuresString, bearerTokenConnectionName, message("iam.name"))
-                )
-                    .yesText(message("credentials.switch.confirmation.yes"))
-                    .noText(message("credentials.switch.confirmation.no"))
-                    .doNotAsk(object : com.intellij.openapi.ui.DoNotAskOption.Adapter() {
-                        override fun rememberChoice(isSelected: Boolean, exitCode: Int) {
-                            if (isSelected && exitCode == DialogWrapper.OK_EXIT_CODE) {
-                                doNotPromptForPinning = true
+        shouldPinConnections.let { shouldPinConnections ->
+            if (shouldPinConnections == null) {
+                val bearerTokenConnectionName = bearerTokenConnectionString(oldConnection, newConnection)
+
+                computeOnEdt(ModalityState.defaultModalityState()) {
+                    MessageDialogBuilder.yesNo(
+                        message("credentials.switch.confirmation.title", featuresString, bearerTokenConnectionName),
+                        message("credentials.switch.confirmation.comment", featuresString, bearerTokenConnectionName, message("iam.name"))
+                    )
+                        .yesText(message("credentials.switch.confirmation.yes"))
+                        .noText(message("credentials.switch.confirmation.no"))
+                        .doNotAsk(object : com.intellij.openapi.ui.DoNotAskOption.Adapter() {
+                            override fun rememberChoice(isSelected: Boolean, exitCode: Int) {
+                                if (isSelected) {
+                                    this@DefaultConnectionPinningManager.shouldPinConnections = exitCode == DialogWrapper.OK_EXIT_CODE
+                                }
+                            }
+                        })
+                        .icon(AllIcons.General.QuestionDialog)
+                        .help(HelpIds.EXPLORER_CREDS_HELP.id)
+                        .ask(project).apply {
+                            if (this) {
+                                UiTelemetry.click(project, "connection_multiple_auths_yes")
+                            } else {
+                                UiTelemetry.click(project, "connection_multiple_auths_no")
                             }
                         }
-                    })
-                    .icon(AllIcons.General.QuestionDialog)
-                    .help(HelpIds.EXPLORER_CREDS_HELP.id)
-                    .ask(project).apply {
-                        if (this) {
-                            UiTelemetry.click(project, "connection_multiple_auths_yes")
-                        } else {
-                            UiTelemetry.click(project, "connection_multiple_auths_no")
-                        }
-                    }
+                }
+            } else {
+                shouldPinConnections
             }
-        } else {
-            false
         }
 
     private fun bearerTokenConnectionString(oldConnection: ToolkitConnection?, newConnection: ToolkitConnection): String {
@@ -167,6 +171,6 @@ class DefaultConnectionPinningManager :
 }
 
 data class ConnectionPinningManagerState(
-    var doNotPromptForPinning: Boolean = false,
+    var shouldPinConnections: Boolean? = null,
     var pinnedConnections: Map<String, String> = emptyMap()
 )

jetbrains-core/tst/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManagerTest.kt

@@ -18,6 +18,7 @@ import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
+import org.mockito.kotlin.verifyNoInteractions
 import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.AwsCredentialConnection
@@ -188,4 +189,43 @@ class ConnectionPinningManagerTest {
 
         assertThat(sut.getPinnedConnection(feature)).isEqualTo(connection)
     }
+
+    @Test
+    fun `respects pinning prompt = yes`() {
+        val connection = mock<AwsBearerTokenConnection>() {
+            on { id } doReturn "connId"
+        }
+        val dialogMock = mock<TestDialog>()
+        TestDialogManager.setTestDialog(dialogMock)
+
+        sut.shouldPinConnections = true
+        assertThat(sut.showDialogIfNeeded(connection, connection, "feature")).isTrue()
+        verifyNoInteractions(dialogMock)
+    }
+
+    @Test
+    fun `respects pinning prompt = no`() {
+        val connection = mock<AwsBearerTokenConnection>() {
+            on { id } doReturn "connId"
+        }
+        val dialogMock = mock<TestDialog>()
+        TestDialogManager.setTestDialog(dialogMock)
+
+        sut.shouldPinConnections = false
+        assertThat(sut.showDialogIfNeeded(connection, connection, "feature")).isFalse()
+        verifyNoInteractions(dialogMock)
+    }
+
+    @Test
+    fun `prompts for pinning`() {
+        val connection = mock<AwsBearerTokenConnection>() {
+            on { id } doReturn "connId"
+        }
+
+        sut.shouldPinConnections = null
+        TestDialogManager.setTestDialog(TestDialog.YES)
+        assertThat(sut.showDialogIfNeeded(connection, connection, "feature")).isTrue()
+        TestDialogManager.setTestDialog(TestDialog.NO)
+        assertThat(sut.showDialogIfNeeded(connection, connection, "feature")).isFalse()
+    }
 }