Setting customization status bar to override if present from AB

Author: spfink

.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"
+}

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

@@ -162,10 +162,11 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
         // A/B case
         val customizationArnFromAB = CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()
+        val customizationNameFromAB = CodeWhispererFeatureConfigService.getInstance().getCustomizationOverrideName()
         if (customizationArnFromAB.isEmpty()) return result
         return CodeWhispererCustomization(
             arn = customizationArnFromAB,
-            name = result?.name.orEmpty(),
+            name = customizationNameFromAB,
             description = result?.description
         )
     }

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

@@ -80,7 +80,7 @@ class CodeWhispererFeatureConfigServiceTest {
                 listOf(
                     FeatureEvaluation.builder()
                         .feature(CodeWhispererFeatureConfigService.CUSTOMIZATION_ARN_OVERRIDE_NAME)
-                        .variation("customizationARN")
+                        .variation("customization-name")
                         .value(FeatureValue.fromStringValue("test arn"))
                         .build()
                 )
@@ -125,6 +125,7 @@ class CodeWhispererFeatureConfigServiceTest {
             assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()).isEqualTo("")
         } else {
             assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationArnOverride()).isEqualTo("test arn")
+            assertThat(CodeWhispererFeatureConfigService.getInstance().getCustomizationOverrideName()).isEqualTo("customization-name")
         }
     }
 

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

@@ -120,8 +120,9 @@ class CodeWhispererModelConfiguratorTest {
 
         abManager.stub {
             on { getCustomizationArnOverride() }.thenReturn("bar")
+            on { getCustomizationOverrideName() }.thenReturn("foo")
         }
-        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "customization_1", "description_1"))
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "foo", "description_1"))
     }
 
     @Test

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt

@@ -107,20 +107,25 @@ class CodeWhispererFeatureConfigService {
     // 5) Specify the corresponding type value getter for the `FeatureValue` class. For example,
     // if the return type is Long, then the corresponding type value getter is `longValue()`.
     // 6) Add a test case for this feature.
-    fun getTestFeature(): String = getFeatureValueForKey(TEST_FEATURE_NAME).stringValue()
+    fun getTestFeature(): String = getFeatureNameForKey(TEST_FEATURE_NAME).stringValue()
 
-    fun getCustomizationArnOverride(): String = getFeatureValueForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
+    fun getCustomizationArnOverride(): String = getFeatureNameForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME).stringValue()
 
-    fun getNewAutoTriggerUX(): Boolean = getFeatureValueForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
+    fun getCustomizationOverrideName(): String = getFeatureVariationForKey(CUSTOMIZATION_ARN_OVERRIDE_NAME)
 
-    fun getInlineCompletion(): Boolean = getFeatureValueForKey(INLINE_COMPLETION).stringValue() == "TREATMENT"
+    fun getNewAutoTriggerUX(): Boolean = getFeatureNameForKey(NEW_AUTO_TRIGGER_UX).stringValue() == "TREATMENT"
+
+    fun getInlineCompletion(): Boolean = getFeatureNameForKey(INLINE_COMPLETION).stringValue() == "TREATMENT"
 
     // 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 =
+    private fun getFeatureNameForKey(name: String): FeatureValue =
         featureConfigs[name]?.value ?: FEATURE_DEFINITIONS[name]?.value
             ?: FeatureValue.builder().boolValue(false).build()
 
+    private fun getFeatureVariationForKey(name: String): String =
+        featureConfigs[name]?.variation.orEmpty()
+
     private fun connection(project: Project) =
         ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())