Add tests for #4468 (#4469)

Author: rli

buildSrc/src/main/kotlin/toolkit-intellij-subplugin.gradle.kts

@@ -122,7 +122,6 @@ tasks.withType<Test>().all {
     val jetbrainsCoreTestResources = project(":plugin-toolkit:jetbrains-core").projectDir.resolve("tst-resources")
     systemProperty("idea.log.config.properties.file", jetbrainsCoreTestResources.resolve("toolkit-test-log.properties"))
     systemProperty("org.gradle.project.ideProfileName", ideProfile.name)
-    systemProperty("aws.dev.useDAG", true)
 
     jvmArgs(openedPackages)
 }

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt

@@ -3,14 +3,18 @@
 
 package software.aws.toolkits.jetbrains.core.credentials.sso
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.testFramework.ApplicationRule
+import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.RuleChain
+import com.intellij.testFramework.replaceService
 import kotlinx.coroutines.runBlocking
 import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
+import org.junit.jupiter.api.assertThrows
 import org.mockito.kotlin.KStubbing
 import org.mockito.kotlin.any
 import org.mockito.kotlin.eq
@@ -33,6 +37,7 @@ import software.amazon.awssdk.services.ssooidc.model.StartDeviceAuthorizationRes
 import software.aws.toolkits.core.region.aRegionId
 import software.aws.toolkits.core.utils.delegateMock
 import software.aws.toolkits.core.utils.test.aString
+import software.aws.toolkits.jetbrains.core.credentials.sso.pkce.ToolkitOAuthService
 import software.aws.toolkits.jetbrains.utils.rules.SsoLoginCallbackProviderRule
 import java.time.Clock
 import java.time.Duration
@@ -53,11 +58,12 @@ class SsoAccessTokenProviderTest {
     private lateinit var ssoCache: SsoCache
 
     private val applicationRule = ApplicationRule()
+    private val disposableRule = DisposableRule()
     private val ssoCallbackRule = SsoLoginCallbackProviderRule()
 
     @JvmField
     @Rule
-    val ruleChain = RuleChain(applicationRule, ssoCallbackRule)
+    val ruleChain = RuleChain(applicationRule, ssoCallbackRule, disposableRule)
 
     @Before
     fun setUp() {
@@ -163,6 +169,58 @@ class SsoAccessTokenProviderTest {
         verify(ssoCache).saveAccessToken(ssoUrl, accessToken)
     }
 
+    @Test
+    fun `initiates authorizatation_grant registration when scopes are requested in a commercial region`() {
+        val sut = SsoAccessTokenProvider(ssoUrl, "us-east-1", ssoCache, ssoOidcClient, scopes = listOf("dummy:scope"), clock = clock)
+        setupCacheStub(returnValue = null)
+
+        val oauth = mock<ToolkitOAuthService>()
+        ApplicationManager.getApplication().replaceService(ToolkitOAuthService::class.java, oauth, disposableRule.disposable)
+
+        ssoOidcClient.stub {
+            on(
+                ssoOidcClient.registerClient(any<RegisterClientRequest>())
+            ).thenReturn(
+                RegisterClientResponse.builder()
+                    .clientId(clientId)
+                    .clientSecret(clientSecret)
+                    .clientSecretExpiresAt(clock.instant().plusSeconds(180).toEpochMilli())
+                    .build()
+            )
+        }
+
+        // flow is not completely stubbed out
+        assertThrows<Exception> { sut.accessToken() }
+
+        verify(ssoCache).saveClientRegistration(any<PKCEClientRegistrationCacheKey>(), any())
+    }
+
+    @Test
+    fun `initiates device code registration when scopes are requested in a non-commercial region`() {
+        val sut = SsoAccessTokenProvider(ssoUrl, "us-gov-east-1", ssoCache, ssoOidcClient, scopes = listOf("dummy:scope"), clock = clock)
+        setupCacheStub(returnValue = null)
+
+        val oauth = mock<ToolkitOAuthService>()
+        ApplicationManager.getApplication().replaceService(ToolkitOAuthService::class.java, oauth, disposableRule.disposable)
+
+        ssoOidcClient.stub {
+            on(
+                ssoOidcClient.registerClient(any<RegisterClientRequest>())
+            ).thenReturn(
+                RegisterClientResponse.builder()
+                    .clientId(clientId)
+                    .clientSecret(clientSecret)
+                    .clientSecretExpiresAt(clock.instant().plusSeconds(180).toEpochMilli())
+                    .build()
+            )
+        }
+
+        // flow is not completely stubbed out
+        assertThrows<Exception> { sut.accessToken() }
+
+        verify(ssoCache).saveClientRegistration(any<DeviceAuthorizationClientRegistrationCacheKey>(), any())
+    }
+
     @Test
     fun getAccessTokenWithoutCachesMultiplePolls() {
         val expirationClientRegistration = clock.instant().plusSeconds(120)