Skip to content

Commit

Permalink
Merge pull request #39565 from drewtul/bug-bedrockagent-kms-key
Browse files Browse the repository at this point in the history 
Bug Fix bug where customer_encryption_key_arn was not passed in on update
  • Loading branch information
@ewbankkit
ewbankkit authored Oct 2, 2024
2 parents ebc475d + b4f4891 commit 4774b44
Show file tree
Hide file tree
Showing 3 changed files with 166 additions and 9 deletions.
7 changes: 7 additions & 0 deletions .changelog/39565.txt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
```release-note:bug
resource/aws_bedrockagent_agent: Fix "Provider produced inconsistent result after apply" error on update due to `customer_encryption_key_arn` not being passed during update
```

```release-note:bug
resource/aws_bedrockagent_agent: Fix "Provider produced inconsistent result after apply" error on update due to `prompt_override_configuration` not being passed when not modified
```
23 changes: 14 additions & 9 deletions internal/service/bedrockagent/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,20 +291,23 @@ func (r *agentResource) Update(ctx context.Context, request resource.UpdateReque
conn := r.Meta().BedrockAgentClient(ctx)

if !new.AgentName.Equal(old.AgentName) ||
!new.AgentResourceRoleARN.Equal(old.AgentResourceRoleARN) ||
!new.CustomerEncryptionKeyARN.Equal(old.CustomerEncryptionKeyARN) ||
!new.Description.Equal(old.Description) ||
!new.Instruction.Equal(old.Instruction) ||
!new.IdleSessionTTLInSeconds.Equal(old.IdleSessionTTLInSeconds) ||
!new.FoundationModel.Equal(old.FoundationModel) ||
!new.GuardrailConfiguration.Equal(old.GuardrailConfiguration) ||
!new.PromptOverrideConfiguration.Equal(old.PromptOverrideConfiguration) {
input := &bedrockagent.UpdateAgentInput{
AgentId: fwflex.StringFromFramework(ctx, new.AgentID),
AgentName: fwflex.StringFromFramework(ctx, new.AgentName),
AgentResourceRoleArn: fwflex.StringFromFramework(ctx, new.AgentResourceRoleARN),
Description: fwflex.StringFromFramework(ctx, new.Description),
FoundationModel: fwflex.StringFromFramework(ctx, new.FoundationModel),
IdleSessionTTLInSeconds: fwflex.Int32FromFramework(ctx, new.IdleSessionTTLInSeconds),
Instruction: fwflex.StringFromFramework(ctx, new.Instruction),
AgentId: fwflex.StringFromFramework(ctx, new.AgentID),
AgentName: fwflex.StringFromFramework(ctx, new.AgentName),
AgentResourceRoleArn: fwflex.StringFromFramework(ctx, new.AgentResourceRoleARN),
CustomerEncryptionKeyArn: fwflex.StringFromFramework(ctx, new.CustomerEncryptionKeyARN),
Description: fwflex.StringFromFramework(ctx, new.Description),
FoundationModel: fwflex.StringFromFramework(ctx, new.FoundationModel),
IdleSessionTTLInSeconds: fwflex.Int32FromFramework(ctx, new.IdleSessionTTLInSeconds),
Instruction: fwflex.StringFromFramework(ctx, new.Instruction),
}

if !new.CustomerEncryptionKeyARN.Equal(old.CustomerEncryptionKeyARN) {
Expand All @@ -321,14 +324,16 @@ func (r *agentResource) Update(ctx context.Context, request resource.UpdateReque
input.GuardrailConfiguration = guardrailConfiguration
}

if !new.PromptOverrideConfiguration.Equal(old.PromptOverrideConfiguration) {
if !new.PromptOverrideConfiguration.IsNull() {
promptOverrideConfiguration := &awstypes.PromptOverrideConfiguration{}
response.Diagnostics.Append(fwflex.Expand(ctx, new.PromptOverrideConfiguration, promptOverrideConfiguration)...)
if response.Diagnostics.HasError() {
return
}

input.PromptOverrideConfiguration = promptOverrideConfiguration
if len(promptOverrideConfiguration.PromptConfigurations) > 0 {
input.PromptOverrideConfiguration = promptOverrideConfiguration
}
}

_, err := conn.UpdateAgent(ctx, input)
Expand Down
145 changes: 145 additions & 0 deletions internal/service/bedrockagent/agent_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,48 @@ func TestAccBedrockAgentAgent_singlePrompt(t *testing.T) {
})
}

func TestAccBedrockAgentAgent_singlePromptUpdate(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
resourceName := "aws_bedrockagent_agent.test"
var v awstypes.Agent

resource.ParallelTest(t, resource.TestCase{
PreCheck: func() { acctest.PreCheck(ctx, t); acctest.PreCheckPartitionHasService(t, names.BedrockEndpointID) },
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
CheckDestroy: testAccCheckAgentDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentConfig_singlePromptUpdate(rName, "anthropic.claude-v2", "basic claude", "500"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
Config: testAccAgentConfig_singlePromptUpdate(rName, "anthropic.claude-v2", "basic claude", "501"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "skip_resource_in_use_check", acctest.CtTrue),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
},
})
}

func TestAccBedrockAgentAgent_addPrompt(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
Expand Down Expand Up @@ -343,6 +385,52 @@ func TestAccBedrockAgentAgent_tags(t *testing.T) {
})
}

func TestAccBedrockAgentAgent_kms(t *testing.T) {
ctx := acctest.Context(t)
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
resourceName := "aws_bedrockagent_agent.test"
var v awstypes.Agent

resource.ParallelTest(t, resource.TestCase{
PreCheck: func() { acctest.PreCheck(ctx, t); acctest.PreCheckPartitionHasService(t, names.BedrockEndpointID) },
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
CheckDestroy: testAccCheckAgentDestroy(ctx),
Steps: []resource.TestStep{
{
Config: testAccAgentConfig_kms(rName, "anthropic.claude-v2", "basic claude", "500"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "guardrail_configuration.#", acctest.Ct0),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "prepare_agent", acctest.CtTrue),
resource.TestCheckResourceAttrPair(resourceName, "customer_encryption_key_arn", "aws_kms_key.test_agent", names.AttrARN),
),
},
{
ResourceName: resourceName,
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"skip_resource_in_use_check"},
},
{
Config: testAccAgentConfig_kms(rName, "anthropic.claude-v2", "basic claude", "501"),
Check: resource.ComposeAggregateTestCheckFunc(
testAccCheckAgentExists(ctx, resourceName, &v),
resource.TestCheckResourceAttr(resourceName, "agent_name", rName),
resource.TestCheckResourceAttr(resourceName, "guardrail_configuration.#", acctest.Ct0),
resource.TestCheckResourceAttr(resourceName, "prompt_override_configuration.#", acctest.Ct1),
resource.TestCheckResourceAttr(resourceName, names.AttrDescription, "basic claude"),
resource.TestCheckResourceAttr(resourceName, "prepare_agent", acctest.CtTrue),
resource.TestCheckResourceAttrPair(resourceName, "customer_encryption_key_arn", "aws_kms_key.test_agent", names.AttrARN),
),
},
},
})
}

func testAccCheckAgentDestroy(ctx context.Context) resource.TestCheckFunc {
return func(s *terraform.State) error {
conn := acctest.Provider.Meta().(*conns.AWSClient).BedrockAgentClient(ctx)
Expand Down Expand Up @@ -651,6 +739,43 @@ resource "aws_bedrockagent_agent" "test" {
`, rName, model, desc))
}

func testAccAgentConfig_singlePromptUpdate(rName, model, desc, timeout string) string {
return acctest.ConfigCompose(testAccAgent_base(rName, model), fmt.Sprintf(`
resource "aws_bedrockagent_agent" "test" {
agent_name = %[1]q
agent_resource_role_arn = aws_iam_role.test_agent.arn
description = %[3]q
idle_session_ttl_in_seconds = %[4]s
instruction = file("${path.module}/test-fixtures/instruction.txt")
foundation_model = %[2]q
skip_resource_in_use_check = true
prompt_override_configuration {
override_lambda = null
prompt_configurations = [
{
base_prompt_template = file("${path.module}/test-fixtures/post-processing.txt")
inference_configuration = [
{
max_length = 2048
stop_sequences = ["Human:"]
temperature = 0
top_k = 250
top_p = 1
},
]
parser_mode = "DEFAULT"
prompt_creation_mode = "OVERRIDDEN"
prompt_state = "DISABLED"
prompt_type = "POST_PROCESSING"
},
]
}
}
`, rName, model, desc, timeout))
}

func testAccAgentConfig_guardrail_noConfig(rName, model string) string {
return acctest.ConfigCompose(testAccAgent_base(rName, model), testAccAgent_guardrail(rName), fmt.Sprintf(`
resource "aws_bedrockagent_agent" "test" {
Expand Down Expand Up @@ -678,3 +803,23 @@ resource "aws_bedrockagent_agent" "test" {
}
`, rName, model, guardrailVersion))
}

func testAccAgentConfig_kms(rName, model, description, timeout string) string {
return acctest.ConfigCompose(testAccAgent_base(rName, model), fmt.Sprintf(`
resource "aws_bedrockagent_agent" "test" {
agent_name = %[1]q
agent_resource_role_arn = aws_iam_role.test_agent.arn
customer_encryption_key_arn = aws_kms_key.test_agent.arn
description = %[3]q
idle_session_ttl_in_seconds = %[4]s
instruction = file("${path.module}/test-fixtures/instruction.txt")
foundation_model = %[2]q
}
resource "aws_kms_key" "test_agent" {
description = "Agent test key"
enable_key_rotation = true
deletion_window_in_days = 7
}
`, rName, model, description, timeout))
}

0 comments on commit 4774b44

Please sign in to comment.