diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitConnectionManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitConnectionManager.kt index ba13fda2a7..c191a3364e 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitConnectionManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitConnectionManager.kt @@ -9,6 +9,7 @@ import com.intellij.openapi.components.PersistentStateComponent import com.intellij.openapi.components.State import com.intellij.openapi.components.Storage import com.intellij.openapi.project.Project +import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection import software.aws.toolkits.jetbrains.core.credentials.pinning.ConnectionPinningManager import software.aws.toolkits.jetbrains.core.credentials.pinning.FeatureWithPinnedConnection import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener @@ -30,6 +31,8 @@ class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStat ) } private val project: Project? + private var keepSecondaryConnection: Boolean = true + constructor(project: Project) { this.project = project } @@ -65,22 +68,21 @@ class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStat } return connection?.let { - if (feature.supportsConnectionType(it)) { + if (feature.supportsConnectionType(it) && keepSecondaryConnection) { return it } - null } ?: defaultConnection?.let { - if (feature.supportsConnectionType(it)) { + if (feature.supportsConnectionType(it) && keepSecondaryConnection) { return it } - null } } override fun getState() = ToolkitConnectionManagerState( - connection?.id + connection?.id, + keepSecondaryConnection ) override fun loadState(state: ToolkitConnectionManagerState) { @@ -94,6 +96,7 @@ class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStat } connection = ToolkitAuthManager.getInstance().getConnection(activeConnectionIdWithRegion) } + keepSecondaryConnection = state.keepSecondaryConnection } @Synchronized @@ -108,7 +111,7 @@ class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStat if (oldConnection != null && newConnection != null && pinningManager != null) { val featuresToPin = mutableListOf() FeatureWithPinnedConnection.EP_NAME.forEachExtensionSafe { - if (!pinningManager.isFeaturePinned(it) && it.supportsConnectionType(oldConnection) && !it.supportsConnectionType(newConnection)) { + if (!pinningManager.isFeaturePinned(it) && (it.supportsConnectionType(oldConnection) && !it.supportsConnectionType(newConnection))) { featuresToPin.add(it) } } @@ -121,8 +124,13 @@ class DefaultToolkitConnectionManager : ToolkitConnectionManager, PersistentStat ApplicationManager.getApplication().messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC).activeConnectionChanged(connection) } } + + override fun setKeepSecondaryConnection(keepSecondaryConnection: Boolean) { + this.keepSecondaryConnection = keepSecondaryConnection + } } data class ToolkitConnectionManagerState( - var activeConnectionId: String? = null + var activeConnectionId: String? = null, + var keepSecondaryConnection: Boolean = true ) diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAddConnectionDialog.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAddConnectionDialog.kt index b9bfc97768..fec0a02c65 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAddConnectionDialog.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAddConnectionDialog.kt @@ -160,6 +160,7 @@ open class ToolkitAddConnectionDialog( } loginSso(project, startUrl, region, scopes) + } } catch (e: Exception) { val message = when (e) { diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt index 2bb6d22a4d..99caece9dd 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt @@ -93,10 +93,13 @@ interface ToolkitConnectionManager : Disposable { fun switchConnection(connection: ToolkitConnection?) + fun setKeepSecondaryConnection(keepSecondaryConnection: Boolean) + override fun dispose() {} companion object { fun getInstance(project: Project?) = project?.let { it.service() } ?: service() + fun getInstance() = service() } } @@ -164,7 +167,6 @@ fun loginSso(project: Project?, startUrl: String, region: String = DEFAULT_SSO_R allScopes.toSet().toList() ) ) - try { reauthConnection(project, connection) } catch (e: Exception) { diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManager.kt index 430a5afb0d..a24c89a959 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/pinning/ConnectionPinningManager.kt @@ -18,6 +18,7 @@ import org.jetbrains.annotations.TestOnly import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection +import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.sono.isSono import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener import software.aws.toolkits.jetbrains.core.help.HelpIds @@ -131,6 +132,8 @@ class DefaultConnectionPinningManager(private val project: Project) : val oldConnectionDisplayName = connectionString(oldConnection) val newConnectionDisplayName = connectionString(newConnection) + val connectionManager = ToolkitConnectionManager.getInstance() + MessageDialogBuilder.yesNo( message("credentials.switch.confirmation.title", featuresString, oldConnectionDisplayName), message("credentials.switch.confirmation.comment", featuresString, oldConnectionDisplayName, newConnectionDisplayName) @@ -152,6 +155,7 @@ class DefaultConnectionPinningManager(private val project: Project) : } else { UiTelemetry.click(project, "connection_multiple_auths_no") } + connectionManager.setKeepSecondaryConnection(this) } } else { false diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt index ac0265c731..0959e256d5 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/core/credentials/sono/SonoCredentialManager.kt @@ -33,7 +33,7 @@ class SonoCredentialManager { } internal fun provider() = ( - ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeCatalystConnection.getInstance()) + ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeCatalystConnection.getInstance()) as? AwsBearerTokenConnection ) ?.getConnectionSettings() @@ -69,6 +69,7 @@ class SonoCredentialManager { val provider = provider() return when (provider?.state()) { null -> runUnderProgressIfNeeded(null, message("credentials.sono.login.pending"), true) { + ToolkitConnectionManager.getInstance().setKeepSecondaryConnection(true) loginSso(project, SONO_URL, requestedScopes = CODECATALYST_SCOPES) } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererWhatIsAction.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererWhatIsAction.kt index ea43e903d4..095498dece 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererWhatIsAction.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererWhatIsAction.kt @@ -21,9 +21,7 @@ class CodeWhispererWhatIsAction : ), DumbAware { override fun update(e: AnActionEvent) { - e.project?.let { - e.presentation.isEnabledAndVisible = isCodeWhispererEnabled(it) - } + e.presentation.isEnabledAndVisible = isCodeWhispererEnabled() } override fun actionPerformed(e: AnActionEvent) { diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt index 332cbfe6f7..f9cab872dc 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanManager.kt @@ -127,7 +127,7 @@ class CodeWhispererCodeScanManager(val project: Project) { * Triggers a code scan and displays results in the new tab in problems view panel. */ fun runCodeScan() { - if (!isCodeWhispererEnabled(project)) return + if (!isCodeWhispererEnabled()) return // Return if a scan is already in progress. if (isCodeScanInProgress.getAndSet(true)) return @@ -165,7 +165,7 @@ class CodeWhispererCodeScanManager(val project: Project) { val startTime = Instant.now().toEpochMilli() var codeScanResponseContext = defaultCodeScanResponseContext() var getProjectSize: Deferred = async { null } - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance()) + val connection = ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeWhispererConnection.getInstance()) try { val file = FileEditorManager.getInstance(project).selectedEditor?.file ?: noFileOpenError() diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/actions/CodeWhispererCodeScanRunAction.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/actions/CodeWhispererCodeScanRunAction.kt index 9951559e5a..6013855d0e 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/actions/CodeWhispererCodeScanRunAction.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/actions/CodeWhispererCodeScanRunAction.kt @@ -17,7 +17,7 @@ class CodeWhispererCodeScanRunAction : DumbAwareAction( ) { override fun update(event: AnActionEvent) { val project = event.project ?: return - event.presentation.isEnabledAndVisible = isCodeWhispererEnabled(project) + event.presentation.isEnabledAndVisible = isCodeWhispererEnabled() val scanManager = CodeWhispererCodeScanManager.getInstance(project) event.presentation.icon = scanManager.getRunActionButtonIcon() } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index eea42fe2ee..eb34d9ec26 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -66,7 +66,7 @@ interface CodeWhispererClientAdaptor : Disposable { fun getInstance(project: Project): CodeWhispererClientAdaptor = project.service() private fun shouldUseSigv4Client(project: Project) = - CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Accountless + CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType() == CodeWhispererLoginType.Accountless } } @@ -158,7 +158,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW ApplicationManager.getApplication().messageBus.syncPublisher(CREDENTIALS_CHANGED) .providerRemoved(oldProviderIdToRemove) - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance()) + val connection = ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeWhispererConnection.getInstance()) connection as? AwsBearerTokenConnection ?: error("$connection is not a bearer token connection") return AwsClientManager.getInstance().getClient(connection.getConnectionSettings()) } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt index fe9a29d2b6..989bbdd9fd 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt @@ -25,7 +25,7 @@ class CodeWhispererEditorListener : EditorFactoryListener { editor.document.addDocumentListener( object : BulkAwareDocumentListener { override fun documentChanged(event: DocumentEvent) { - if (!isCodeWhispererEnabled(project)) return + if (!isCodeWhispererEnabled()) return CodeWhispererInvocationStatus.getInstance().documentChanged() CodeWhispererCodeCoverageTracker.getInstance(language).apply { activateTrackerIfNotActive() diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererExplorerActionManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererExplorerActionManager.kt index cafa479e8d..6967f50464 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererExplorerActionManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererExplorerActionManager.kt @@ -45,7 +45,7 @@ class CodeWhispererExplorerActionManager : PersistentStateComponent CodeWhispererLoginType.Accountless - isAccessTokenExpired(project) || isRefreshTokenExpired(project) -> CodeWhispererLoginType.Expired + isAccessTokenExpired() || isRefreshTokenExpired() -> CodeWhispererLoginType.Expired else -> { - val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance()) + val conn = ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeWhispererConnection.getInstance()) if (conn != null) { if (conn.isSono()) { CodeWhispererLoginType.Sono @@ -196,10 +196,10 @@ interface CodeWhispererActivationChangedListener { fun activationChanged(value: Boolean) {} } -fun isCodeWhispererEnabled(project: Project) = with(CodeWhispererExplorerActionManager.getInstance()) { - checkActiveCodeWhispererConnectionType(project) != CodeWhispererLoginType.Logout +fun isCodeWhispererEnabled() = with(CodeWhispererExplorerActionManager.getInstance()) { + checkActiveCodeWhispererConnectionType() != CodeWhispererLoginType.Logout } -fun isCodeWhispererExpired(project: Project) = with(CodeWhispererExplorerActionManager.getInstance()) { - checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Expired +fun isCodeWhispererExpired() = with(CodeWhispererExplorerActionManager.getInstance()) { + checkActiveCodeWhispererConnectionType() == CodeWhispererLoginType.Expired } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererServiceNode.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererServiceNode.kt index 700601fc45..eef5490d58 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererServiceNode.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/CodeWhispererServiceNode.kt @@ -81,7 +81,7 @@ class CodeWhispererServiceNode( } val manager = CodeWhispererExplorerActionManager.getInstance() - val activeConnectionType = manager.checkActiveCodeWhispererConnectionType(project) + val activeConnectionType = manager.checkActiveCodeWhispererConnectionType() return when (activeConnectionType) { CodeWhispererLoginType.Logout -> listOf(whatIsCodeWhispererNode, getStartedCodeWhispererNode) @@ -106,7 +106,7 @@ class CodeWhispererServiceNode( return } - val connectionType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) + val connectionType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType() when (connectionType) { CodeWhispererLoginType.Expired -> { presentation.addText(message("codewhisperer.explorer.root_node.login_type.expired"), SimpleTextAttributes.GRAY_ATTRIBUTES) diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/nodes/GetStartedNode.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/nodes/GetStartedNode.kt index 7ac532aea2..78d23c254e 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/nodes/GetStartedNode.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/explorer/nodes/GetStartedNode.kt @@ -46,10 +46,11 @@ class GetStartedNode(nodeProject: Project) : CodeWhispererActionNode( */ private fun enableCodeWhisperer(project: Project) { val explorerActionManager = CodeWhispererExplorerActionManager.getInstance() - val connectionManager = ToolkitConnectionManager.getInstance(project) + val connectionManager = ToolkitConnectionManager.getInstance() connectionManager.activeConnectionForFeature(CodeWhispererConnection.getInstance())?.let { project.refreshDevToolTree() } ?: run { + connectionManager.setKeepSecondaryConnection(true) runInEdt { // Start from scratch if no active connection if (CodeWhispererLoginDialog(project).showAndGet()) { @@ -58,7 +59,7 @@ class GetStartedNode(nodeProject: Project) : CodeWhispererActionNode( } } - if (isCodeWhispererEnabled(project)) { + if (isCodeWhispererEnabled()) { StartupActivity.POST_STARTUP_ACTIVITY.findExtension(CodeWhispererProjectStartupActivity::class.java)?.let { it.runActivity(project) } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt index 0763b67158..f7ffb72f02 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt @@ -77,7 +77,7 @@ class CodeWhispererService { latencyContext: LatencyContext ) { val project = editor.project ?: return - if (!isCodeWhispererEnabled(project)) return + if (!isCodeWhispererEnabled()) return latencyContext.credentialFetchingStart = System.nanoTime() @@ -490,7 +490,7 @@ class CodeWhispererService { ): RequestContext { val fileContextInfo = getFileContextInfo(editor, psiFile) val caretPosition = getCaretPosition(editor) - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance()) + val connection = ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeWhispererConnection.getInstance()) return RequestContext(project, editor, triggerTypeInfo, caretPosition, fileContextInfo, connection, latencyContext) } @@ -573,10 +573,8 @@ class CodeWhispererService { } fun canDoInvocation(editor: Editor, type: CodewhispererTriggerType): Boolean { - editor.project?.let { - if (!isCodeWhispererEnabled(it)) { - return false - } + if (!isCodeWhispererEnabled()) { + return false } if (type == CodewhispererTriggerType.AutoTrigger && !CodeWhispererExplorerActionManager.getInstance().isAutoEnabled()) { diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt index b0519bab0d..377672e83a 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt @@ -24,13 +24,13 @@ class CodeWhispererConfigurable(private val project: Project) : get() = CodeWhispererSettings.getInstance() private val isSso: Boolean - get() = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.SSO + get() = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType() == CodeWhispererLoginType.SSO override fun getId() = "aws.codewhisperer" override fun createPanel() = panel { val connect = project.messageBus.connect(disposable ?: error("disposable wasn't initialized by framework")) - val invoke = isCodeWhispererEnabled(project) + val invoke = isCodeWhispererEnabled() // TODO: can we remove message bus subscribe and solely use visible(boolean) / enabled(boolean), consider multi project cases row { @@ -42,7 +42,7 @@ class CodeWhispererConfigurable(private val project: Project) : ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - visible(!isCodeWhispererEnabled(project)) + visible(!isCodeWhispererEnabled()) } } ) @@ -55,7 +55,7 @@ class CodeWhispererConfigurable(private val project: Project) : ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - enabled(isCodeWhispererEnabled(project) && !isSso) + enabled(isCodeWhispererEnabled() && !isSso) } } ) @@ -70,7 +70,7 @@ class CodeWhispererConfigurable(private val project: Project) : ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - enabled(isCodeWhispererEnabled(project)) + enabled(isCodeWhispererEnabled()) } } ) @@ -85,7 +85,7 @@ class CodeWhispererConfigurable(private val project: Project) : ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - enabled(isCodeWhispererEnabled(project) && !isSso) + enabled(isCodeWhispererEnabled() && !isSso) } } ) diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt index 50513b21fc..2a00ceaeaa 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt @@ -38,7 +38,7 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware { if (!ApplicationManager.getApplication().isUnitTestMode) { CodeWhispererStatusBarManager.getInstance(project).updateWidget() } - if (!isCodeWhispererEnabled(project)) return + if (!isCodeWhispererEnabled()) return if (runOnce) return // Reconnect CodeWhisperer on startup @@ -56,7 +56,7 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware { } private fun showAccountlessNotificationIfNeeded(project: Project) { - if (CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) == CodeWhispererLoginType.Accountless) { + if (CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType() == CodeWhispererLoginType.Accountless) { // simply show a notification when user login with Accountless, and it's still supported by CodeWhisperer if (!isExpired()) { // don't show warn notification if user selected Don't show again or if notification was shown less than a week ago diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupSettingsListener.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupSettingsListener.kt index 1c38acc18f..6f1eb2d705 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupSettingsListener.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupSettingsListener.kt @@ -37,7 +37,7 @@ class CodeWhispererProjectStartupSettingsListener(private val project: Project) override fun toolWindowShown(toolWindow: ToolWindow) { super.toolWindowShown(toolWindow) if (toolWindow.id != ProblemsView.ID) return - if (!isCodeWhispererEnabled(project)) return + if (!isCodeWhispererEnabled()) return CodeWhispererCodeScanManager.getInstance(project).addCodeScanUI() } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarManager.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarManager.kt index aa910b8081..aa2d27128e 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarManager.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarManager.kt @@ -52,7 +52,7 @@ class CodeWhispererStatusBarManager(private val project: Project) : Disposable { ExtensionPointName("com.intellij.statusBarWidgetFactory").extensionList.find { it.id == CodeWhispererStatusBarWidgetFactory.ID }?.let { - settings.setEnabled(it, isCodeWhispererEnabled(project)) + settings.setEnabled(it, isCodeWhispererEnabled()) widgetsManager.updateWidget(it) } } diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidget.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidget.kt index 51ca16cfc8..cce6c9d8bf 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidget.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidget.kt @@ -55,7 +55,7 @@ class CodeWhispererStatusBarWidget(project: Project) : override fun getClickConsumer(): Consumer? = null override fun getPopupStep(): ListPopup? = - if (isCodeWhispererExpired(project)) { + if (isCodeWhispererExpired()) { JBPopupFactory.getInstance().createConfirmation(message("codewhisperer.statusbar.popup.title"), { reconnectCodeWhisperer(project) }, 0) } else { null @@ -64,7 +64,7 @@ class CodeWhispererStatusBarWidget(project: Project) : override fun getSelectedValue(): String = message("codewhisperer.statusbar.display_name") override fun getIcon(): Icon = - if (isCodeWhispererExpired(project)) { + if (isCodeWhispererExpired()) { AllIcons.General.BalloonWarning } else if (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) { AnimatedIcon.Default() diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidgetFactory.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidgetFactory.kt index 4c36c0bd5c..39c8795f2a 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidgetFactory.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidgetFactory.kt @@ -16,8 +16,7 @@ class CodeWhispererStatusBarWidgetFactory : StatusBarEditorBasedWidgetFactory() override fun getDisplayName(): String = message("codewhisperer.statusbar.display_name") - override fun isAvailable(project: Project): Boolean = - isCodeWhispererEnabled(project) + override fun isAvailable(project: Project): Boolean = isCodeWhispererEnabled() override fun createWidget(project: Project): StatusBarWidget = CodeWhispererStatusBarWidget(project) diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceComponents.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceComponents.kt index 087095f3ae..c1f385f2a1 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceComponents.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceComponents.kt @@ -65,22 +65,22 @@ class CodeWhispererCodeReferenceComponents(private val project: Project) { }.asCodeReferencePanelFont() init { - repaint(project) + repaint() // set the reference panel text different for SSO users vs AWS Builder ID / Accless users project.messageBus.connect().subscribe( ToolkitConnectionManagerListener.TOPIC, object : ToolkitConnectionManagerListener { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { - repaint(project) + repaint() } } ) } // TODO: figure out how to have a different view for SSO user in a cleaner way, maybe have 2 sets of components stored in [ReferenceManager]? - private fun repaint(project: Project) { - val loginType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType(project) + private fun repaint() { + val loginType = CodeWhispererExplorerActionManager.getInstance().checkActiveCodeWhispererConnectionType() settingsLabelPrefixText as JLabel settingsLabelLink as ActionLink if (loginType == CodeWhispererLoginType.SSO) { diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceToolWindowFactory.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceToolWindowFactory.kt index 3619bc3381..65dacdc6fe 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceToolWindowFactory.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/toolwindow/CodeWhispererCodeReferenceToolWindowFactory.kt @@ -21,7 +21,7 @@ class CodeWhispererCodeReferenceToolWindowFactory : ToolWindowFactory, DumbAware toolWindow.contentManager.addContent(toolWindowContent) } - override fun shouldBeAvailable(project: Project): Boolean = isCodeWhispererEnabled(project) + override fun shouldBeAvailable(project: Project): Boolean = isCodeWhispererEnabled() companion object { const val id = "aws.codewhisperer.codereference" diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt index af4be7d704..b525a94f94 100644 --- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt +++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererUtil.kt @@ -94,14 +94,14 @@ object CodeWhispererUtil { listOf(CodeWhispererSsoLearnMoreAction(), ConnectWithAwsToContinueActionError(), DoNotShowAgainActionError()) ) - fun isAccessTokenExpired(project: Project): Boolean { - val tokenProvider = tokenProvider(project) ?: return false + fun isAccessTokenExpired(): Boolean { + val tokenProvider = tokenProvider() ?: return false val state = tokenProvider.state() return state == BearerTokenAuthState.NEEDS_REFRESH } - fun isRefreshTokenExpired(project: Project): Boolean { - val tokenProvider = tokenProvider(project) ?: return false + fun isRefreshTokenExpired(): Boolean { + val tokenProvider = tokenProvider() ?: return false val state = tokenProvider.state() return state == BearerTokenAuthState.NOT_AUTHENTICATED } @@ -113,9 +113,9 @@ object CodeWhispererUtil { // Return true if need to re-auth, false otherwise fun promptReAuth(project: Project): Boolean { if (CodeWhispererService.hasReAuthPromptBeenShown()) return false - if (!isCodeWhispererExpired(project)) return false - val tokenProvider = tokenProvider(project) ?: return false - return maybeReauthProviderIfNeeded(project, tokenProvider) { + if (!isCodeWhispererExpired()) return false + val tokenProvider = tokenProvider() ?: return false + return maybeReauthProviderIfNeeded(null, tokenProvider) { runInEdt { notifyConnectionExpiredRequestReauth(project) CodeWhispererService.markReAuthPromptShown() @@ -130,7 +130,7 @@ object CodeWhispererUtil { notifyError( message("toolkit.sso_expire.dialog.title"), message("toolkit.sso_expire.dialog_message"), - project, + null, listOf( NotificationAction.create(message("toolkit.sso_expire.dialog.yes_button")) { _, notification -> reconnectCodeWhisperer(project) @@ -150,9 +150,9 @@ object CodeWhispererUtil { return connection.startUrl } - private fun tokenProvider(project: Project) = ( + private fun tokenProvider() = ( ToolkitConnectionManager - .getInstance(project) + .getInstance() .activeConnectionForFeature(CodeWhispererConnection.getInstance()) as? AwsBearerTokenConnection ) ?.getConnectionSettings() @@ -160,7 +160,7 @@ object CodeWhispererUtil { ?.delegate as? BearerTokenProvider fun reconnectCodeWhisperer(project: Project) { - val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(CodeWhispererConnection.getInstance()) + val connection = ToolkitConnectionManager.getInstance().activeConnectionForFeature(CodeWhispererConnection.getInstance()) if (connection !is BearerSsoConnection) return ApplicationManager.getApplication().executeOnPooledThread { getConnectionStartUrl(connection)?.let { startUrl -> diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt index 0c33baa599..07a4d8a3d4 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt @@ -3,6 +3,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer +import com.intellij.openapi.application.ApplicationManager import com.intellij.testFramework.DisposableRule import com.intellij.testFramework.ProjectRule import com.intellij.testFramework.RuleChain @@ -109,7 +110,7 @@ class CodeWhispererClientAdaptorTest { activeConnectionForFeature(any()) } doReturn authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), emptyList())) as AwsBearerTokenConnection } - projectRule.project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable) + ApplicationManager.getApplication().replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable) } @Test diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt index 0aee66e1c4..b4f3014ff9 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt @@ -111,7 +111,7 @@ internal abstract class CodeWhispererCodeCoverageTrackerTestBase(myProjectRule: telemetryServiceSpy = spy(TestTelemetryService(batcher = batcher)) exploreActionManagerMock = mock { - on { checkActiveCodeWhispererConnectionType(any()) } doReturn CodeWhispererLoginType.Sono + on { checkActiveCodeWhispererConnectionType() } doReturn CodeWhispererLoginType.Sono } ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreActionManagerMock, disposableRule.disposable) diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererExplorerActionManagerTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererExplorerActionManagerTest.kt index 9cdc820b35..147c54a9fa 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererExplorerActionManagerTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererExplorerActionManagerTest.kt @@ -4,7 +4,6 @@ package software.aws.toolkits.jetbrains.services.codewhisperer import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.project.Project import com.intellij.testFramework.ApplicationRule import com.intellij.testFramework.DisposableRule import com.intellij.testFramework.ProjectRule @@ -51,16 +50,13 @@ class CodeWhispererExplorerActionManagerTest { val mockClientManager = MockClientManagerRule() private lateinit var mockManager: CodeWhispererExplorerActionManager - private lateinit var project: Project private lateinit var connectionManager: ToolkitConnectionManager @Before fun setup() { mockClientManager.create() - project = projectRule.project connectionManager = mock() - - project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable) + ApplicationManager.getApplication().replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable) } /** @@ -71,7 +67,7 @@ class CodeWhispererExplorerActionManagerTest { mockManager = spy() whenever(connectionManager.activeConnectionForFeature(any())).thenReturn(null) - val actual = mockManager.checkActiveCodeWhispererConnectionType(project) + val actual = mockManager.checkActiveCodeWhispererConnectionType() assertThat(actual).isEqualTo(CodeWhispererLoginType.Logout) } @@ -85,7 +81,7 @@ class CodeWhispererExplorerActionManagerTest { } ) - val actual = mockManager.checkActiveCodeWhispererConnectionType(project) + val actual = mockManager.checkActiveCodeWhispererConnectionType() assertThat(actual).isEqualTo(CodeWhispererLoginType.Accountless) } @@ -124,17 +120,17 @@ class CodeWhispererExplorerActionManagerTest { mockManager = mock() ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, mockManager, disposableRule.disposable) - whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Logout) - assertThat(isCodeWhispererEnabled(project)).isFalse + whenever(mockManager.checkActiveCodeWhispererConnectionType()).thenReturn(CodeWhispererLoginType.Logout) + assertThat(isCodeWhispererEnabled()).isFalse - whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Accountless) - assertThat(isCodeWhispererEnabled(project)).isTrue + whenever(mockManager.checkActiveCodeWhispererConnectionType()).thenReturn(CodeWhispererLoginType.Accountless) + assertThat(isCodeWhispererEnabled()).isTrue - whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.Sono) - assertThat(isCodeWhispererEnabled(project)).isTrue + whenever(mockManager.checkActiveCodeWhispererConnectionType()).thenReturn(CodeWhispererLoginType.Sono) + assertThat(isCodeWhispererEnabled()).isTrue - whenever(mockManager.checkActiveCodeWhispererConnectionType(project)).thenReturn(CodeWhispererLoginType.SSO) - assertThat(isCodeWhispererEnabled(project)).isTrue + whenever(mockManager.checkActiveCodeWhispererConnectionType()).thenReturn(CodeWhispererLoginType.SSO) + assertThat(isCodeWhispererEnabled()).isTrue } private fun assertLoginType(startUrl: String, expectedType: CodeWhispererLoginType) { @@ -144,7 +140,7 @@ class CodeWhispererExplorerActionManagerTest { whenever(conn.startUrl).thenReturn(startUrl) whenever(conn.getConnectionSettings()).thenReturn(null) - val actual = mockManager.checkActiveCodeWhispererConnectionType(project) + val actual = mockManager.checkActiveCodeWhispererConnectionType() assertThat(actual).isEqualTo(expectedType) } } diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt index 0dbd9fa9fe..ba91ccc859 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt @@ -87,8 +87,8 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() { @Test fun `when isCodeWhispererEnabled is false, user not able to trigger CodeWhisperer manually`() { - whenever(stateManager.checkActiveCodeWhispererConnectionType(projectRule.project)).thenReturn(CodeWhispererLoginType.Logout) - assertThat(isCodeWhispererEnabled(projectRule.project)).isFalse + whenever(stateManager.checkActiveCodeWhispererConnectionType()).thenReturn(CodeWhispererLoginType.Logout) + assertThat(isCodeWhispererEnabled()).isFalse invokeCodeWhispererService() verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any()) } diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererStatusBarManagerTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererStatusBarManagerTest.kt index 3b6c504d7c..245a77ebe8 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererStatusBarManagerTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererStatusBarManagerTest.kt @@ -15,7 +15,6 @@ import com.intellij.testFramework.replaceService import org.junit.Before import org.junit.Rule import org.junit.Test -import org.mockito.kotlin.any import org.mockito.kotlin.eq import org.mockito.kotlin.mock import org.mockito.kotlin.verify @@ -96,6 +95,6 @@ class CodeWhispererStatusBarManagerTest { } private fun configureConnection(loginType: CodeWhispererLoginType) { - whenever(explorerActionManager.checkActiveCodeWhispererConnectionType(any())).thenReturn(loginType) + whenever(explorerActionManager.checkActiveCodeWhispererConnectionType()).thenReturn(loginType) } } diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt index b90ebb69eb..5d3ec9897a 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt @@ -436,7 +436,7 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { @Test fun `test codePercentage tracker will not be activated if CWSPR terms of service is not accepted`() { val exploreManagerMock = mock { - on { checkActiveCodeWhispererConnectionType(projectRule.project) } doReturn CodeWhispererLoginType.Logout + on { checkActiveCodeWhispererConnectionType() } doReturn CodeWhispererLoginType.Logout } ApplicationManager.getApplication().replaceService(CodeWhispererExplorerActionManager::class.java, exploreManagerMock, disposableRule.disposable) val project = projectRule.project diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt index bf6a08d3b0..15efc257a3 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt @@ -188,7 +188,7 @@ open class CodeWhispererTestBase { fun mockCodeWhispererEnabledStatus(enabled: Boolean) { stateManager.stub { onGeneric { - checkActiveCodeWhispererConnectionType(any()) + checkActiveCodeWhispererConnectionType() } doAnswer { if (enabled) CodeWhispererLoginType.Sono else CodeWhispererLoginType.Logout } diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt index 333a28f2fe..61ddadb547 100644 --- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt +++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanTestBase.kt @@ -99,7 +99,7 @@ open class CodeWhispererCodeScanTestBase(projectRule: CodeInsightTestFixtureRule ApplicationManager.getApplication().replaceService( CodeWhispererExplorerActionManager::class.java, mock { - on { checkActiveCodeWhispererConnectionType(any()) } doReturn CodeWhispererLoginType.Accountless + on { checkActiveCodeWhispererConnectionType() } doReturn CodeWhispererLoginType.Accountless }, disposableRule.disposable )