Fixing test.

builtin/providers/aws/resource_aws_ami_launch_permission_test.go

@@ -2,17 +2,19 @@ package aws
 
 import (
 	"fmt"
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/sts"
 	r "github.com/hashicorp/terraform/helper/resource"
 	"github.com/hashicorp/terraform/terraform"
 	"testing"
 )
 
+type testLaunchPermission struct {
+	ImageID   string
+	AccountID string
+}
+
 func TestAccAWSAMILaunchPermission_Basic(t *testing.T) {
-	var amiId string
-	var accountId *string
+	tlp := &testLaunchPermission{}
 	r.Test(t, r.TestCase{
 		PreCheck: func() {
 			testAccPreCheck(t)
@@ -21,19 +23,19 @@ func TestAccAWSAMILaunchPermission_Basic(t *testing.T) {
 			if err != nil {
 				t.Fatalf("could not initialize ami launch permission test: %s", err)
 			}
-			*accountId = *res.Account
+			tlp.AccountID = *res.Account
 		},
 		Providers: testAccProviders,
 		Steps: []r.TestStep{
 			r.TestStep{
-				Config: testAccAWSAMILaunchPermissionConfig(*accountId),
+				Config: testAccAWSAMILaunchPermissionConfig(tlp),
 				Check: r.ComposeTestCheckFunc(
-					testCheckResourceGetAttr("aws_ami_from_instance", "id", &amiId),
-					testAccAWSAMILaunchPermissionExists(amiId, *accountId),
+					testCheckResourceGetAttr("aws_ami_from_instance", "id", &tlp.AccountID),
+					testAccAWSAMILaunchPermissionExists(tlp),
 				),
 			},
 		},
-		CheckDestroy: testAccAWSAMILaunchPermissionDestroyed(),
+		CheckDestroy: testAccAWSAMILaunchPermissionDestroyed(tlp),
 	})
 }
 
@@ -55,38 +57,31 @@ func testCheckResourceGetAttr(name, key string, value *string) r.TestCheckFunc {
 	}
 }
 
-func testAccAWSAMILaunchPermissionExists(amiId string, accountId string) r.TestCheckFunc {
+func testAccAWSAMILaunchPermissionExists(tlp *testLaunchPermission) r.TestCheckFunc {
 	return func(s *terraform.State) error {
 		conn := testAccProvider.Meta().(*AWSClient).ec2conn
-		res, err := conn.DescribeImageAttribute(&ec2.DescribeImageAttributeInput{
-			ImageId:   aws.String(amiId),
-			Attribute: aws.String("launchPermission"),
-		})
-		if err != nil {
+		if has, err := hasLaunchPermission(conn, tlp.ImageID, tlp.AccountID); err != nil {
 			return err
+		} else if !has {
+			return fmt.Errorf("launch permission does not exist for '%s' on '%s'", tlp.AccountID, tlp.ImageID)
 		}
-		for _, lp := range res.LaunchPermissions {
-			if *lp.UserId == accountId {
-				return nil
-			}
-		}
-		return fmt.Errorf("launch permission does not exist for '%s' on '%s'", accountId, amiId)
+		return nil
 	}
 }
 
-func testAccAWSAMILaunchPermissionDestroyed() r.TestCheckFunc {
+func testAccAWSAMILaunchPermissionDestroyed(tlp *testLaunchPermission) r.TestCheckFunc {
 	return func(s *terraform.State) error {
 		conn := testAccProvider.Meta().(*AWSClient).ec2conn
-
-		conn.DescribeImageAttribute(&ec2.DescribeImageAttributeInput{
-			ImageId: aws.String(""),
-		})
-
+		if has, err := hasLaunchPermission(conn, tlp.ImageID, tlp.AccountID); err != nil {
+			return err
+		} else if has {
+			return fmt.Errorf("launch permission still exists for '%s' on '%s'", tlp.AccountID, tlp.ImageID)
+		}
 		return nil
 	}
 }
 
-func testAccAWSAMILaunchPermissionConfig(accountId string) string {
+func testAccAWSAMILaunchPermissionConfig(tlp *testLaunchPermission) string {
 	return fmt.Sprintf(`
 provider "aws" {
 	region = "us-east-1"
@@ -137,5 +132,5 @@ resource "aws_ami_launch_permission" "self-test" {
     image_id   = "${aws_ami_from_instance.test.id}"
     account_id = "%s"
 }
-`, accountId)
+`, tlp.AccountID)
 }