@@ -7,8 +7,10 @@ import com.intellij.openapi.application.ApplicationManager
77import com.intellij.testFramework.ApplicationRule
88import com.intellij.testFramework.DisposableRule
99import com.intellij.testFramework.ProjectRule
10+ import com.intellij.testFramework.registerServiceInstance
1011import com.intellij.testFramework.replaceService
1112import com.intellij.util.xmlb.XmlSerializer
13+ import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
1214import org.assertj.core.api.Assertions.assertThat
1315import org.jdom.output.XMLOutputter
1416import org.junit.Before
@@ -40,10 +42,14 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
4042import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
4143import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
4244import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
45+ import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
46+ import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
4347import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4448import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomizationState
4549import software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator
4650import software.aws.toolkits.jetbrains.utils.xmlElement
51+ import java.util.concurrent.CountDownLatch
52+ import java.util.concurrent.TimeUnit
4753import kotlin.reflect.full.memberProperties
4854import kotlin.reflect.jvm.isAccessible
4955
@@ -75,6 +81,7 @@ class CodeWhispererModelConfiguratorTest {
7581 private lateinit var sut: DefaultCodeWhispererModelConfigurator
7682 private lateinit var mockClient: CodeWhispererRuntimeClient
7783 private lateinit var abManager: CodeWhispererFeatureConfigService
84+ private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
7885
7986 @Before
8087 fun setup () {
@@ -83,7 +90,11 @@ class CodeWhispererModelConfiguratorTest {
8390 regionProvider.addRegion(Region .US_EAST_1 )
8491 regionProvider.addRegion(Region .US_EAST_2 )
8592
86- sut = DefaultCodeWhispererModelConfigurator ()
93+ sut = spy(CodeWhispererModelConfigurator .getInstance() as DefaultCodeWhispererModelConfigurator ).also { spyInstance ->
94+ ApplicationManager .getApplication().replaceService(
95+ DefaultCodeWhispererModelConfigurator ::class .java, spyInstance, disposableRule.disposable
96+ )
97+ }
8798
8899 (ToolkitConnectionManager .getInstance(projectRule.project) as DefaultToolkitConnectionManager ).loadState(ToolkitConnectionManagerState ())
89100 mockClient.stub {
@@ -110,6 +121,9 @@ class CodeWhispererModelConfiguratorTest {
110121 abManager,
111122 disposableRule.disposable
112123 )
124+
125+ mockClintAdaptor = mock()
126+ projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor ::class .java, mockClintAdaptor)
113127 }
114128
115129 @Test
@@ -550,4 +564,48 @@ class CodeWhispererModelConfiguratorTest {
550564 assertThat(actual.previousAvailableCustomizations).hasSize(1 )
551565 assertThat(actual.previousAvailableCustomizations[" fake-sso-url" ]).isEqualTo(listOf (" arn_1" , " arn_2" , " arn_3" ))
552566 }
567+
568+ @Test
569+ fun `profile switch should keep using existing customization if new list still contains that arn` () {
570+ val ssoConn = spy(LegacyManagedBearerSsoConnection (region = " us-east-1" , startUrl = " url 1" , scopes = Q_SCOPES ))
571+ ToolkitConnectionManager .getInstance(projectRule.project).switchConnection(ssoConn)
572+ val oldCustomization = CodeWhispererCustomization (" oldArn" , " oldName" , " oldDescription" )
573+ sut.switchCustomization(projectRule.project, oldCustomization)
574+
575+ assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
576+
577+ val fakeCustomizations = listOf (
578+ CodeWhispererCustomization (" oldArn" , " oldName" , " oldDescription" )
579+ )
580+ mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
581+
582+ ApplicationManager .getApplication().messageBus
583+ .syncPublisher(QRegionProfileSelectedListener .TOPIC )
584+ .onProfileSelected(projectRule.project, null )
585+
586+ assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
587+ }
588+
589+ @Test
590+ fun `profile switch should invalidate obsolete customization if it's not in the new list` () {
591+ val ssoConn = spy(LegacyManagedBearerSsoConnection (region = " us-east-1" , startUrl = " url 1" , scopes = Q_SCOPES ))
592+ ToolkitConnectionManager .getInstance(projectRule.project).switchConnection(ssoConn)
593+ val oldCustomization = CodeWhispererCustomization (" oldArn" , " oldName" , " oldDescription" )
594+ sut.switchCustomization(projectRule.project, oldCustomization)
595+ assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
596+ val fakeCustomizations = listOf (
597+ CodeWhispererCustomization (" newArn" , " newName" , " newDescription" )
598+ )
599+ mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
600+
601+ val latch = CountDownLatch (1 )
602+
603+ ApplicationManager .getApplication().messageBus
604+ .syncPublisher(QRegionProfileSelectedListener .TOPIC )
605+ .onProfileSelected(projectRule.project, null )
606+
607+ latch.await(2 , TimeUnit .SECONDS )
608+
609+ assertThat(sut.activeCustomization(projectRule.project)).isNull()
610+ }
553611}
0 commit comments