From 09a5ec6e14c0887a32fc3f0c671e7c03ce65fc6d Mon Sep 17 00:00:00 2001 From: Sam Fink Date: Tue, 7 Jan 2025 12:08:54 -0800 Subject: [PATCH] Change logic for customization override to not override if user has manually selected a customization --- ...-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json | 4 +++ .../CodeWhispererModelConfigurator.kt | 27 +++++++++++------- .../CodeWhispererModelConfiguratorTest.kt | 28 ++++++++++++++++--- 3 files changed, 45 insertions(+), 14 deletions(-) create mode 100644 .changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json diff --git a/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json b/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json new file mode 100644 index 0000000000..82086fc5af --- /dev/null +++ b/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json @@ -0,0 +1,4 @@ +{ + "type" : "bugfix", + "description" : "Change logic for customization override to not override if user has manually selected a customization" +} \ No newline at end of file diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt index 3cacc48325..f6535a4bba 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt @@ -157,17 +157,24 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe return@calculateIfIamIdentityCenterConnection customizationUiItems } + /** + * Gets the active customization for a user. If a user has manually selected a customization, + * respect that choice. If a user has not selected a customization, check if they have a customization + * assigned to them via an AB feature. If so, use that customization. + */ override fun activeCustomization(project: Project): CodeWhispererCustomization? { - val result = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] } - - // A/B case - val customizationFeature = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature() - if (customizationFeature == null || customizationFeature.value.stringValue().isEmpty()) return result - return CodeWhispererCustomization( - arn = customizationFeature.value.stringValue(), - name = customizationFeature.variation, - description = result?.description - ) + val selectedCustomization = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] } + + if (selectedCustomization != null) { + return selectedCustomization + } else { + val customizationOverride = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature() + if (customizationOverride == null || customizationOverride.value.stringValue().isEmpty()) return null + return CodeWhispererCustomization( + arn = customizationOverride.value.stringValue(), + name = customizationOverride.variation, + ) + } } override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?) { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt index 169473a97e..d960870240 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt @@ -113,17 +113,37 @@ class CodeWhispererModelConfiguratorTest { } @Test - fun `should override customization arn if there is one under AB test`() { + fun `should not override customization arn if there is one under AB test and manual selection has been made`() { val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES)) ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn) - sut.switchCustomization(projectRule.project, CodeWhispererCustomization("foo", "customization_1", "description_1")) - assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("foo", "customization_1", "description_1")) + sut.switchCustomization(projectRule.project, CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1")) + assertThat(sut.activeCustomization(projectRule.project)) + .isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1")) abManager.stub { on { getCustomizationFeature() }.thenReturn(FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("bar").build())) } - assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "foo", "description_1")) + assertThat(sut.activeCustomization(projectRule.project)) + .isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1")) + } + + @Test + fun `should override customization arn if there is one under AB test and manual selection has not been made`() { + val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES)) + ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn) + + sut.switchCustomization(projectRule.project, CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1")) + assertThat(sut.activeCustomization(projectRule.project)) + .isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1")) + sut.invalidateCustomization("selectedCustomizationArn") + + abManager.stub { + on { getCustomizationFeature() }.thenReturn( + FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("overrideArn").build()) + ) + } + assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("overrideArn", "foo", null)) } @Test