generate separate checkFunc

internal/generate/tags/main.go

@@ -1,9 +1,6 @@
 // Copyright (c) HashiCorp, Inc.
 // SPDX-License-Identifier: MPL-2.0
 
-//go:build generate
-// +build generate
-
 package main
 
 import (

internal/generate/tags/templates/v2/update_tags_body.gtpl

@@ -164,28 +164,8 @@ func {{ .UpdateTagsFunc }}(ctx context.Context, conn {{ .ClientType }}, identifi
 
 	{{ if .WaitForPropagation }}
 	if len(removedTags) > 0 || len(updatedTags) > 0 {
-		checkFunc := func(tags tftags.KeyValueTags) func() (bool, error) {
-			return func() (bool, error) {
-				output, err := listTags(ctx, conn, identifier, optFns...)
-
-				if tfresource.NotFound(err) {
-					return false, nil
-				}
-
-				if err != nil {
-					return false, err
-				}
-
-				if inContext, ok := tftags.FromContext(ctx); ok {
-					tags = tags.IgnoreConfig(inContext.IgnoreConfig)
-					output = output.IgnoreConfig(inContext.IgnoreConfig)
-				}
-
-				return output.Equal(tags), nil
-			}
-		}
-
-		if err := {{ .WaitTagsPropagatedFunc }}(ctx, conn, identifier, newTags, checkFunc(newTags), optFns...); err != nil {
+        check := checkFunc(ctx, conn, newTags, identifier, optFns...)
+		if err := {{ .WaitTagsPropagatedFunc }}(ctx, newTags, check); err != nil {
 			return fmt.Errorf("waiting for resource (%s) tag propagation: %w", identifier, err)
 		}
 	}

internal/generate/tags/templates/v2/update_tags_for_resource_body.gtpl

@@ -164,9 +164,9 @@ func {{ .UpdateTagsForResourceFunc }}(ctx context.Context, conn {{ .ClientType }
 
 	{{ if .WaitForPropagation }}
 	if len(removedTags) > 0 || len(updatedTags) > 0 {
-		checkFunc := func(tags tftags.KeyValueTags) func() (bool, error) {
+		checkFunc := func(ctx context.Context, conn {{ .ClientType }}, tags tftags.KeyValueTags, id string, optFns ...func(*{{ .AWSService }}.Options)) func() (bool, error) {
 			return func() (bool, error) {
-				output, err := listTags(ctx, conn, identifier, optFns...)
+				output, err := listTags(ctx, conn, id, optFns...)
 
 				if tfresource.NotFound(err) {
 					return false, nil
@@ -185,7 +185,7 @@ func {{ .UpdateTagsForResourceFunc }}(ctx context.Context, conn {{ .ClientType }
 			}
 		}
 
-		if err := {{ .WaitTagsPropagatedFunc }}(ctx, conn, identifier, newTags, checkFunc(newTags), optFns...); err != nil {
+		if err := {{ .WaitTagsPropagatedFunc }}(ctx, newTags, checkFunc(ctx, conn, newTags, identifier, optFns...)); err != nil {
 			return fmt.Errorf("waiting for resource (%s) tag propagation: %w", identifier, err)
 		}
 	}

internal/generate/tags/templates/v2/wait_tags_propagated_body.gtpl

@@ -1,7 +1,7 @@
 // {{ .WaitTagsPropagatedFunc }} waits for {{ .ServicePackage }} service tags to be propagated.
 // The identifier is typically the Amazon Resource Name (ARN), although
 // it may also be a different identifier depending on the service.
-func {{ .WaitTagsPropagatedFunc }}(ctx context.Context, conn {{ .ClientType }}, id string, tags tftags.KeyValueTags, checkFunc func() (bool, error), optFns ...func(*{{ .AWSService }}.Options)) error {
+func {{ .WaitTagsPropagatedFunc }}(ctx context.Context, tags tftags.KeyValueTags, checkFunc func() (bool, error)) error {
 	tflog.Debug(ctx, "Waiting for tag propagation", map[string]any{
 		names.AttrTags: tags,
 	})
@@ -23,3 +23,29 @@ func {{ .WaitTagsPropagatedFunc }}(ctx context.Context, conn {{ .ClientType }},
 
 	return tfresource.WaitUntil(ctx, {{ .WaitTimeout }}, checkFunc, opts)
 }
+
+// checkFunc returns a function that checks if the tags are propagated.
+func checkFunc(ctx context.Context, conn {{ .ClientType }}, tags tftags.KeyValueTags, id string, optFns ...func(*{{ .AWSService }}.Options)) func() (bool, error) {
+    return func() (bool, error) {
+        output, err := listTags(ctx, conn, id, optFns...)
+
+        if tfresource.NotFound(err) {
+            return false, nil
+        }
+
+        if err != nil {
+            return false, err
+        }
+
+        if inContext, ok := tftags.FromContext(ctx); ok {
+            tags = tags.IgnoreConfig(inContext.IgnoreConfig)
+            output = output.IgnoreConfig(inContext.IgnoreConfig)
+        }
+
+        {{- if .UpdateTagsForResource }}
+        return output.ContainsAll(tags), nil
+        {{- else }}
+        return output.Equal(tags), nil
+        {{- end }}
+    }
+}

internal/service/kms/external_key.go

@@ -204,7 +204,8 @@ func resourceExternalKeyCreate(ctx context.Context, d *schema.ResourceData, meta
 	}
 
 	if tags := KeyValueTags(ctx, getTagsIn(ctx)); len(tags) > 0 {
-		if err := waitTagsPropagated(ctx, conn, d.Id(), tags); err != nil {
+		check := checkFunc(ctx, conn, tags, d.Id())
+		if err := waitTagsPropagated(ctx, tags, check); err != nil {
 			return sdkdiag.AppendErrorf(diags, "waiting for KMS External Key (%s) tag update: %s", d.Id(), err)
 		}
 	}

internal/service/kms/key.go

@@ -218,7 +218,8 @@ func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, meta interfa
 	}
 
 	if tags := KeyValueTags(ctx, getTagsIn(ctx)); len(tags) > 0 {
-		if err := waitTagsPropagated(ctx, conn, d.Id(), tags); err != nil {
+		check := checkFunc(ctx, conn, tags, d.Id())
+		if err := waitTagsPropagated(ctx, tags, check); err != nil {
 			return sdkdiag.AppendErrorf(diags, "waiting for KMS Key (%s) tag update: %s", d.Id(), err)
 		}
 	}

internal/service/kms/replica_external_key.go

@@ -199,7 +199,8 @@ func resourceReplicaExternalKeyCreate(ctx context.Context, d *schema.ResourceDat
 	}
 
 	if tags := KeyValueTags(ctx, getTagsIn(ctx)); len(tags) > 0 {
-		if err := waitTagsPropagated(ctx, conn, d.Id(), tags); err != nil {
+		check := checkFunc(ctx, conn, tags, d.Id())
+		if err := waitTagsPropagated(ctx, tags, check); err != nil {
 			return sdkdiag.AppendErrorf(diags, "waiting for KMS Replica External Key (%s) tag update: %s", d.Id(), err)
 		}
 	}

internal/service/kms/replica_key.go

@@ -172,7 +172,8 @@ func resourceReplicaKeyCreate(ctx context.Context, d *schema.ResourceData, meta
 	}
 
 	if tags := KeyValueTags(ctx, getTagsIn(ctx)); len(tags) > 0 {
-		if err := waitTagsPropagated(ctx, conn, d.Id(), tags); err != nil {
+		check := checkFunc(ctx, conn, tags, d.Id())
+		if err := waitTagsPropagated(ctx, tags, check); err != nil {
 			return sdkdiag.AppendErrorf(diags, "waiting for KMS Replica Key (%s) tag update: %s", d.Id(), err)
 		}
 	}