Merge main into feature/q-mega

aws · Nov 20, 2024 · 5d9b765 · 5d9b765
2 parents 5505137 + 7fd65de
commit 5d9b765
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 12 deletions.
diff --git a/.changes/next-release/feature-625c1f2c-3dfd-47cb-8907-408041a3642c.json b/.changes/next-release/feature-625c1f2c-3dfd-47cb-8907-408041a3642c.json
@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Uses AB variation as the name for overriden customizations"
+}
diff --git a/...toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/...toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt
@@ -161,11 +161,11 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         val result = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] }
 
         // A/B case
-        val customizationArnFromAB = CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()
-        if (customizationArnFromAB.isEmpty()) return result
+        val customizationFeature = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()
+        if (customizationFeature == null || customizationFeature.value.stringValue().isEmpty()) return result
         return CodeWhispererCustomization(
-            arn = customizationArnFromAB,
-            name = result?.name.orEmpty(),
+            arn = customizationFeature.value.stringValue(),
+            name = customizationFeature.variation,
             description = result?.description
         )
     }

diff --git a/...re/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt b/...re/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt
@@ -14,6 +14,7 @@ import org.junit.Rule
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
+import org.mockito.kotlin.isNull
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.stub
 import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
@@ -80,7 +81,7 @@ class CodeWhispererFeatureConfigServiceTest {
                 listOf(
                     FeatureEvaluation.builder()
                         .feature(CodeWhispererFeatureConfigService.CUSTOMIZATION_ARN_OVERRIDE_NAME)
-                        .variation("customizationARN")
+                        .variation("customization-name")
                         .value(FeatureValue.fromStringValue("test arn"))
                         .build()
                 )
@@ -122,9 +123,10 @@ class CodeWhispererFeatureConfigServiceTest {
         }
 
         if (!isIdc || !isInListAvailableCustomizations) {
-            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()).isEqualTo("")
+            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()).isNull()
         } else {
-            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()).isEqualTo("test arn")
+            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.value?.stringValue()).isEqualTo("test arn")
+            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.variation).isEqualTo("customization-name")
         }
     }
 

diff --git a/...tware/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt b/...tware/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt
@@ -22,6 +22,7 @@ import org.mockito.kotlin.stub
 import software.amazon.awssdk.regions.Region
 import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
 import software.amazon.awssdk.services.codewhispererruntime.model.Customization
+import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue
 import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsResponse
 import software.amazon.awssdk.services.ssooidc.SsoOidcClient
@@ -38,6 +39,7 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
 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.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomizationState
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator
@@ -100,7 +102,7 @@ class CodeWhispererModelConfiguratorTest {
         }
 
         abManager = mock {
-            on { getCustomizationArnOverride() }.thenReturn("")
+            on { getCustomizationFeature() }.thenReturn(null)
         }
 
         ApplicationManager.getApplication().replaceService(
@@ -119,9 +121,9 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("foo", "customization_1", "description_1"))
 
         abManager.stub {
-            on { getCustomizationArnOverride() }.thenReturn("bar")
+            on { getCustomizationFeature() }.thenReturn(FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("bar").build()))
         }
-        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "customization_1", "description_1"))
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "foo", "description_1"))
     }
 
     @Test

diff --git a/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/...src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt
@@ -109,7 +109,7 @@ class CodeWhispererFeatureConfigService {
     // 6) Add a test case for this feature.
     fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue()
 
-    fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
+    fun getCustomizationFeature(): FeatureContext? = getFeature(CUSTOMIZATION_ARN_OVERRIDE_NAME)
 
     fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
 
@@ -118,9 +118,12 @@ class CodeWhispererFeatureConfigService {
     // Get the feature value for the given key.
     // In case of a misconfiguration, it will return a default feature value of Boolean false.
     private fun getFeatureValueForKey(name: String): FeatureValue =
-        featureConfigs[name]?.value ?: FEATURE_DEFINITIONS[name]?.value
+        getFeature(name)?.value ?: FEATURE_DEFINITIONS[name]?.value
             ?: FeatureValue.builder().boolValue(false).build()
 
+    // Gets the feature context for a given feature name.
+    private fun getFeature(name: String): FeatureContext? = featureConfigs[name]
+
     private fun connection(project: Project) =
         ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())