fix: fix invalid result object after apply issue

internal/service/xray/resource_policy.go

@@ -112,6 +112,50 @@ func (r *resourcePolicyResource) Create(ctx context.Context, req resource.Create
 	resp.Diagnostics.Append(resp.State.Set(ctx, plan)...)
 }
 
+func (r *resourcePolicyResource) Update(ctx context.Context, req resource.UpdateRequest, resp *resource.UpdateResponse) {
+	conn := r.Meta().XRayClient(ctx)
+
+	var plan resourcePolicyResourceModel
+	resp.Diagnostics.Append(req.Plan.Get(ctx, &plan)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	in := xray.PutResourcePolicyInput{
+		PolicyDocument: plan.PolicyDocument.ValueStringPointer(),
+		PolicyName:     plan.PolicyName.ValueStringPointer(),
+	}
+	resp.Diagnostics.Append(fwflex.Expand(ctx, plan, &in)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	_, err := conn.PutResourcePolicy(ctx, &in)
+	if err != nil {
+		resp.Diagnostics.AddError(
+			create.ProblemStandardMessage(names.XRay, create.ErrActionUpdating, ResNameResourcePolicy, plan.PolicyName.String(), err),
+			err.Error(),
+		)
+		return
+	}
+
+	out, err := findResourcePolicyByName(ctx, conn, plan.PolicyName.ValueString())
+	if err != nil {
+		resp.Diagnostics.AddError(
+			create.ProblemStandardMessage(names.XRay, create.ErrActionSetting, ResNameResourcePolicy, plan.PolicyName.String(), err),
+			err.Error(),
+		)
+		return
+	}
+
+	resp.Diagnostics.Append(fwflex.Flatten(ctx, out, &plan)...)
+	if resp.Diagnostics.HasError() {
+		return
+	}
+
+	resp.Diagnostics.Append(resp.State.Set(ctx, plan)...)
+}
+
 func (r *resourcePolicyResource) Read(ctx context.Context, req resource.ReadRequest, resp *resource.ReadResponse) {
 	conn := r.Meta().XRayClient(ctx)
 

internal/service/xray/resource_policy_test.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/xray/types"
 	sdkacctest "github.com/hashicorp/terraform-plugin-testing/helper/acctest"
 	"github.com/hashicorp/terraform-plugin-testing/helper/resource"
@@ -44,6 +45,36 @@ func TestAccXRayResourcePolicy_basic(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "policy_revision_id", "1"),
 				),
 			},
+		},
+	})
+}
+
+func TestAccXRayResourcePolicy_policyDocument(t *testing.T) {
+	ctx := acctest.Context(t)
+	var resourcepolicy types.ResourcePolicy
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+	resourceName := "aws_xray_resource_policy.test"
+	policyDocument1 := `{"Version":"2012-10-17","Statement":[{"Sid":"AllowXRayAccess","Effect":"Allow","Principal":{"AWS":"*"},"Action":["xray:*","xray:PutResourcePolicy"],"Resource":"*"}]}`
+	policyDocument2 := `{"Version":"2012-10-17","Statement":[{"Sid":"AllowXRayAccessUpdated","Effect":"Allow","Principal":{"AWS":"*"},"Action":["xray:PutTraceSegments","xray:PutTelemetryRecords"],"Resource":"*"}]}`
+	policyDocument3 := `{"Version":"2012-10-17","Statement":[{"Sid":"AllowXRayAccessFinal","Effect":"Allow","Principal":{"Service":"sns.amazonaws.com"},"Action":"xray:PutTraceSegments","Resource":"*"}]}`
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck: func() {
+			acctest.PreCheck(ctx, t)
+		},
+		ErrorCheck:               acctest.ErrorCheck(t, names.XRayServiceID),
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
+		CheckDestroy:             testAccCheckResourcePolicyDestroy(ctx),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccResourcePolicyConfig_policyDocument(rName, policyDocument1),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckResourcePolicyExists(ctx, resourceName, &resourcepolicy),
+					resource.TestCheckResourceAttr(resourceName, "policy_name", rName),
+					resource.TestCheckResourceAttr(resourceName, "policy_revision_id", "1"),
+					testAccCheckResourcePolicyHasPolicyDocument(ctx, resourceName, policyDocument1),
+				),
+			},
 			{
 				ResourceName:                         resourceName,
 				ImportState:                          true,
@@ -54,6 +85,26 @@ func TestAccXRayResourcePolicy_basic(t *testing.T) {
 					"bypass_policy_lockout_check",
 				},
 			},
+			{
+				Config: testAccResourcePolicyConfig_policyDocument(rName, policyDocument2),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckResourcePolicyExists(ctx, resourceName, &resourcepolicy),
+					resource.TestCheckResourceAttr(resourceName, "policy_name", rName),
+					resource.TestCheckResourceAttr(resourceName, "policy_revision_id", "2"),
+					resource.TestCheckResourceAttrSet(resourceName, names.AttrLastUpdatedTime),
+					testAccCheckResourcePolicyHasPolicyDocument(ctx, resourceName, policyDocument2),
+				),
+			},
+			{
+				Config: testAccResourcePolicyConfig_policyDocument(rName, policyDocument3),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckResourcePolicyExists(ctx, resourceName, &resourcepolicy),
+					resource.TestCheckResourceAttr(resourceName, "policy_name", rName),
+					resource.TestCheckResourceAttr(resourceName, "policy_revision_id", "3"),
+					resource.TestCheckResourceAttrSet(resourceName, names.AttrLastUpdatedTime),
+					testAccCheckResourcePolicyHasPolicyDocument(ctx, resourceName, policyDocument3),
+				),
+			},
 		},
 	})
 }
@@ -156,3 +207,38 @@ resource "aws_xray_resource_policy" "test" {
 }
 `, rName)
 }
+
+func testAccResourcePolicyConfig_policyDocument(rName, policyDocument string) string {
+	return fmt.Sprintf(`
+resource "aws_xray_resource_policy" "test" {
+  policy_name                 = %[1]q
+  policy_document             = %[2]q
+  bypass_policy_lockout_check = true
+}
+`, rName, policyDocument)
+}
+
+func testAccCheckResourcePolicyHasPolicyDocument(ctx context.Context, name string, expectedDocument string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		rs, ok := s.RootModule().Resources[name]
+		if !ok {
+			return create.Error(names.XRay, create.ErrActionCheckingExistence, tfxray.ResNameResourcePolicy, name, errors.New("not found"))
+		}
+
+		conn := acctest.Provider.Meta().(*conns.AWSClient).XRayClient(ctx)
+		output, err := tfxray.FindResourcePolicyByName(ctx, conn, rs.Primary.Attributes["policy_name"])
+		if err != nil {
+			return err
+		}
+
+		if output.PolicyDocument == nil {
+			return fmt.Errorf("policy_document is nil")
+		}
+
+		if aws.ToString(output.PolicyDocument) != expectedDocument {
+			return fmt.Errorf("policy_document mismatch:\nexpected: %s\nactual:   %s", expectedDocument, aws.ToString(output.PolicyDocument))
+		}
+
+		return nil
+	}
+}