add ut

Author: evanliu048

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

@@ -50,6 +50,7 @@ private fun notifyInvalidSelectedCustomization(project: Project) {
 }
 
 private fun notifyNewCustomization(project: Project) {
+    if (ApplicationManager.getApplication().isUnitTestMode) return
     notifyInfo(
         title = message("codewhisperer.custom.dialog.title"),
         content = message("codewhisperer.notification.custom.new_customization"),

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt

@@ -7,8 +7,10 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.testFramework.ApplicationRule
 import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.ProjectRule
+import com.intellij.testFramework.registerServiceInstance
 import com.intellij.testFramework.replaceService
 import com.intellij.util.xmlb.XmlSerializer
+import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import org.assertj.core.api.Assertions.assertThat
 import org.jdom.output.XMLOutputter
 import org.junit.Before
@@ -40,10 +42,14 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
 import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomizationState
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.utils.xmlElement
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import kotlin.reflect.full.memberProperties
 import kotlin.reflect.jvm.isAccessible
 
@@ -75,6 +81,7 @@ class CodeWhispererModelConfiguratorTest {
     private lateinit var sut: DefaultCodeWhispererModelConfigurator
     private lateinit var mockClient: CodeWhispererRuntimeClient
     private lateinit var abManager: CodeWhispererFeatureConfigService
+    private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
 
     @Before
     fun setup() {
@@ -83,7 +90,13 @@ class CodeWhispererModelConfiguratorTest {
         regionProvider.addRegion(Region.US_EAST_1)
         regionProvider.addRegion(Region.US_EAST_2)
 
-        sut = DefaultCodeWhispererModelConfigurator()
+        val original = CodeWhispererModelConfigurator.getInstance()
+        sut = spy(original) as DefaultCodeWhispererModelConfigurator
+        ApplicationManager.getApplication().replaceService(
+            DefaultCodeWhispererModelConfigurator::class.java,
+            sut,
+            disposableRule.disposable
+        )
 
         (ToolkitConnectionManager.getInstance(projectRule.project) as DefaultToolkitConnectionManager).loadState(ToolkitConnectionManagerState())
         mockClient.stub {
@@ -110,6 +123,9 @@ class CodeWhispererModelConfiguratorTest {
             abManager,
             disposableRule.disposable
         )
+
+        mockClintAdaptor = mock()
+        projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor::class.java, mockClintAdaptor)
     }
 
     @Test
@@ -550,4 +566,48 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }
+
+    @Test
+    fun `profile switch should keep using existing customization if new list still contains that arn`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+    }
+
+    @Test
+    fun `profile switch should invalidate obsolete customization if it's not in the new list`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("newArn", "newName", "newDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        val latch = CountDownLatch(1)
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        latch.await(2, TimeUnit.SECONDS)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isNull()
+    }
 }