r/aws_ami: Refactor tagging logic to keyvaluetags package

aws/resource_aws_ami.go

@@ -15,6 +15,7 @@ import (
 	"github.com/hashicorp/terraform-plugin-sdk/helper/hashcode"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
+	"github.com/terraform-providers/terraform-provider-aws/aws/internal/keyvaluetags"
 )
 
 const (
@@ -278,12 +279,18 @@ func resourceAwsAmiCreate(d *schema.ResourceData, meta interface{}) error {
 	id := *res.ImageId
 	d.SetId(id)
 
+	if v := d.Get("tags").(map[string]interface{}); len(v) > 0 {
+		if err := keyvaluetags.Ec2UpdateTags(client, id, nil, v); err != nil {
+			return fmt.Errorf("error adding tags: %s", err)
+		}
+	}
+
 	_, err = resourceAwsAmiWaitForAvailable(d.Timeout(schema.TimeoutCreate), id, client)
 	if err != nil {
 		return err
 	}
 
-	return resourceAwsAmiUpdate(d, meta)
+	return resourceAwsAmiRead(d, meta)
 }
 
 func resourceAwsAmiRead(d *schema.ResourceData, meta interface{}) error {
@@ -404,10 +411,12 @@ func resourceAwsAmiUpdate(d *schema.ResourceData, meta interface{}) error {
 
 	d.Partial(true)
 
-	if err := setTags(client, d); err != nil {
-		return err
-	} else {
-		d.SetPartial("tags")
+	if d.HasChange("tags") {
+		o, n := d.GetChange("tags")
+
+		if err := keyvaluetags.Ec2UpdateTags(client, d.Id(), o, n); err != nil {
+			return fmt.Errorf("error updating AMI (%s) tags: %s", d.Id(), err)
+		}
 	}
 
 	if d.Get("description").(string) != "" {

aws/resource_aws_ami_test.go

@@ -26,7 +26,7 @@ func TestAccAWSAMI_basic(t *testing.T) {
 		CheckDestroy: testAccCheckAmiDestroy,
 		Steps: []resource.TestStep{
 			{
-				Config: testAccAmiConfig_basic(rName),
+				Config: testAccAmiConfig_basic(rName, 8),
 				Check: resource.ComposeTestCheckFunc(
 					testAccCheckAmiExists(resourceName, &ami),
 					resource.TestCheckResourceAttr(resourceName, "ena_support", "true"),
@@ -46,6 +46,53 @@ func TestAccAWSAMI_basic(t *testing.T) {
 	})
 }
 
+func TestAccAWSAMI_tags(t *testing.T) {
+	var ami ec2.Image
+	resourceName := "aws_ami.test"
+	rName := acctest.RandomWithPrefix("tf-acc-test")
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckAmiDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAmiConfigTags1(rName, "key1", "value1", 8),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAmiExists(resourceName, &ami),
+					resource.TestCheckResourceAttr(resourceName, "tags.%", "1"),
+					resource.TestCheckResourceAttr(resourceName, "tags.key1", "value1"),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+				ImportStateVerifyIgnore: []string{
+					"manage_ebs_snapshots",
+				},
+			},
+			{
+				Config: testAccAmiConfigTags2(rName, "key1", "value1updated", "key2", "value2", 8),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAmiExists(resourceName, &ami),
+					resource.TestCheckResourceAttr(resourceName, "tags.%", "2"),
+					resource.TestCheckResourceAttr(resourceName, "tags.key1", "value1updated"),
+					resource.TestCheckResourceAttr(resourceName, "tags.key2", "value2"),
+				),
+			},
+			{
+				Config: testAccAmiConfigTags1(rName, "key2", "value2", 8),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAmiExists(resourceName, &ami),
+					resource.TestCheckResourceAttr(resourceName, "tags.%", "1"),
+					resource.TestCheckResourceAttr(resourceName, "tags.key2", "value2"),
+				),
+			},
+		},
+	})
+}
+
 func TestAccAWSAMI_snapshotSize(t *testing.T) {
 	var ami ec2.Image
 	var bd ec2.BlockDeviceMapping
@@ -66,7 +113,7 @@ func TestAccAWSAMI_snapshotSize(t *testing.T) {
 		CheckDestroy: testAccCheckAmiDestroy,
 		Steps: []resource.TestStep{
 			{
-				Config: testAccAmiConfig_snapshotSize(rName),
+				Config: testAccAmiConfig_basic(rName, 20),
 				Check: resource.ComposeTestCheckFunc(
 					testAccCheckAmiExists(resourceName, &ami),
 					testAccCheckAmiBlockDevice(&ami, &bd, "/dev/sda1"),
@@ -218,30 +265,35 @@ func testAccCheckAmiEbsBlockDevice(bd *ec2.BlockDeviceMapping, ed *ec2.EbsBlockD
 	}
 }
 
-func testAccAmiConfig_basic(rName string) string {
+func testAccAmiConfig_base(size int, rName string) string {
 	return fmt.Sprintf(`
 data "aws_availability_zones" "available" {}
 
 resource "aws_ebs_volume" "foo" {
   availability_zone = "${data.aws_availability_zones.available.names[0]}"
-  size              = 8
+  size              = %d
 
   tags = {
-    Name = "testAccAmiConfig_basic"
+    Name = "%s"
   }
 }
 
 resource "aws_ebs_snapshot" "foo" {
   volume_id = "${aws_ebs_volume.foo.id}"
 
   tags = {
-    Name = "testAccAmiConfig_basic"
+    Name = "%s"
   }
 }
 
+`, size, rName, rName)
+}
+
+func testAccAmiConfig_basic(rName string, size int) string {
+	return testAccAmiConfig_base(size, rName) + fmt.Sprintf(`
 resource "aws_ami" "test" {
   ena_support         = true
-  name                = %q
+  name                = %[1]q
   root_device_name    = "/dev/sda1"
   virtualization_type = "hvm"
 
@@ -253,36 +305,43 @@ resource "aws_ami" "test" {
 `, rName)
 }
 
-func testAccAmiConfig_snapshotSize(rName string) string {
-	return fmt.Sprintf(`
-data "aws_availability_zones" "available" {}
-
-resource "aws_ebs_volume" "foo" {
-  availability_zone = "${data.aws_availability_zones.available.names[0]}"
-  size              = 20
+func testAccAmiConfigTags1(rName, tagKey1, tagValue1 string, size int) string {
+	return testAccAmiConfig_base(size, rName) + fmt.Sprintf(`
+resource "aws_ami" "test" {
+  ena_support         = true
+  name                = %[1]q
+  root_device_name    = "/dev/sda1"
+  virtualization_type = "hvm"
 
-  tags = {
-    Name = "testAccAmiConfig_snapshotSize"
+  ebs_block_device {
+    device_name = "/dev/sda1"
+    snapshot_id = "${aws_ebs_snapshot.foo.id}"
   }
-}
-
-resource "aws_ebs_snapshot" "foo" {
-  volume_id = "${aws_ebs_volume.foo.id}"
 
   tags = {
-    Name = "TestAccAWSAMI_snapshotSize"
+    %[2]q = %[3]q
   }
 }
+`, rName, tagKey1, tagValue1)
+}
 
+func testAccAmiConfigTags2(rName, tagKey1, tagValue1, tagKey2, tagValue2 string, size int) string {
+	return testAccAmiConfig_base(size, rName) + fmt.Sprintf(`
 resource "aws_ami" "test" {
-  name                = %q
+  ena_support         = true
+  name                = %[1]q
   root_device_name    = "/dev/sda1"
   virtualization_type = "hvm"
 
   ebs_block_device {
     device_name = "/dev/sda1"
     snapshot_id = "${aws_ebs_snapshot.foo.id}"
   }
+
+  tags = {
+    %[2]q = %[3]q
+    %[4]q = %[5]q
+  }
 }
-`, rName)
+`, rName, tagKey1, tagValue1, tagKey2, tagValue2)
 }